# Towards Joint Sequence-Structure Generation
of Nucleic Acid and Protein Complexes
with $\mathrm{SE(3)}$-Discrete Diffusion

Alex Morehead , Jeffrey A. Ruffolo, Aadyot Bhatnagar, & Ali Madani
AI Research Lab
Profluent Bio
Berkeley, CA, USA
<EMAIL_ADDRESS>{jruffolo, abhatnagar<EMAIL_ADDRESS>
Alex Morehead was employed by Profluent Bio as a research intern at the time
of this work.

###### Abstract

Generative models of macromolecules carry abundant and impactful implications
for industrial and biomedical efforts in protein engineering. However,
existing methods are currently limited to modeling protein structures or
sequences, independently or jointly, without regard to the interactions that
commonly occur between proteins and other macromolecules. In this work, we
introduce MMDiff, a generative model that jointly designs sequences and
structures of nucleic acid and protein complexes, independently or in complex,
using joint $\mathrm{SE(3)}$-discrete diffusion noise. Such a model has
important implications for emerging areas of macromolecular design including
structure-based transcription factor design and design of noncoding RNA
sequences. We demonstrate the utility of MMDiff through a rigorous new design
benchmark for macromolecular complex generation that we introduce in this
work. Our results demonstrate that MMDiff is able to successfully generate
micro-RNA and single-stranded DNA molecules while being modestly capable of
joint modeling DNA and RNA molecules in interaction with multi-chain protein
complexes. Source code: https://github.com/Profluent-Internships/MMDiff.

##  1 Introduction

The ability to effectively design proteins using computational methods is a
long-sought-after goal of protein engineering research (Lutz, 2010). Such
techniques could enable accelerated development of vaccines targeting novel
viruses, enhanced exploration of the space of designable materials, and the
ability to engineer new protein-based energy resources towards a more
sustainable climate future. Recent progress in computationally designing
protein variants given an initial wildtype sequence has demonstrated the
utility of machine learning in protein engineering (Hie and Yang, 2022). More
recently, deep learning models based on structure prediction networks have
demonstrated the ability to flexibly design novel protein structures with
desired functions in wet-lab experiments (Watson et al., 2023). Nonetheless,
considerably less attention has been directed toward the design of novel
protein-nucleic acid interactions, which carry implications for
computationally engineering new means of transcription factors and intervening
in a host of essential functions within living organisms. Approaching such
design tasks is arguably more challenging, as the number of high-quality
macromolecular complexes currently available in the RCSB Protein Data Bank
(PDB) (Berman et al., 2000) is notably scarce while the shapes of such
complexes can also be much more flexible compared to proteins.

Towards this end, in this work we explore the design of new nucleic acid (NA)
and protein complexes using a generative diffusion process we refer to as
$\mathrm{SE(3)}$-discrete diffusion. In Section 2.2, we discuss the
$\mathrm{SE(3)}$-discrete process which guides deep learning models to jointly
reverse an $\mathrm{SE(3)}$ diffusion process 111$\mathrm{SE(3)}$ refers to
the group of 3D rotations and translations (excluding reflections), whereas
$\mathrm{SE(3)}^{N}$ represents the manifold of $N$ $\mathrm{SE(3)}$ frames.
on rigid body frames in $\mathbb{R}^{3}$ (Yim et al., 2023) and a related
diffusion process on discrete sequence inputs (Chen et al., 2022). In Section
2.3, we introduce our proposed $\mathrm{SE(3)}$-discrete diffusion method for
protein-nucleic acid design, MMDiff, which is to our best knowledge the first
generative diffusion model for the design of macromolecular sequences and
structures beyond proteins. Through our experiments in Section 3, we introduce
a systematic open-source suite of computational benchmarks to assess the
ability of a generative macromolecular design model to produce designable,
diverse, and novel samples, and, under such a framework, highlight the utility
of MMDiff for macromolecular complex generation. Consequently, our
contributions presented in this work pave the way for the development of
future macromolecular generative models.

Figure 1: Our proposed Macromolecular Diffusion Model (MMDiff) jointly designs
macromolecular sequences and structures. A. An overview of MMDiff. For each
nucleic acid residue, a rigid body frame centered at the
$\mathrm{C4^{\prime}}$ atom is constructed. To build such a frame, the
$\mathrm{GramSchmidt}$ algorithm is applied to a residue’s $v_{1}$ and $v_{2}$
vectors, in the process placing its $C3^{\prime}$, $O4^{\prime}$, and
$C5^{\prime}$ atoms with respect to the position of the $C4^{\prime}$ atom.
The positions of all other residue atoms are placed autoregressively according
to a corresponding torsion angle ($\Phi$) predicted by MMDiff. B. An
illustrative example of how MMDiff generates realistic macromolecular samples.
Through the iterative process of denoising $N$ geometric frames and $N$ one-
hot sequence vectors initialized from their respective reference
distributions, MMDiff transitions an initially-random sequence-structure pair
at timestep $\mathrm{T_{F}}$ into a coherent macromolecule at timestep $0$, at
which point each frame and its associated torsion angles are used to construct
the position of each atom.

##  2 Methodology

###  2.1 Preliminaries and Notation

Protein geometry parametrization. Inspired by the recent work of Yim et al.
(2023), we adopt the standard backbone parametrization of AlphaFold 2 for
protein geometry modeling. Succinctly, with such a parametrization, $N_{p}$
protein residues are modeled as $N_{p}$ rigid body frames centered at each
residue’s $\mathrm{C^{\star}_{\alpha}}$ atom
($\mathrm{C^{\star}_{\alpha}}=(0,0,0)$), with each frame representing a
residue’s $\mathrm{N}^{\star},\mathrm{C^{\star}_{\alpha}},\mathrm{C}^{\star}$,
and $\mathrm{O}^{\star}$ (local) atoms, respectively. For a given protein
residue indexed by $n$, to obtain the absolute (global) positions of its main
atoms, one may apply a special Euclidean ($\mathrm{SE(3)}$) transformation
$T_{n}$ as

$[\mathrm{N}_{n},\mathrm{C}_{n},(\mathrm{C_{\alpha}})_{n}]=T_{n}\cdot[\mathrm{N}^{\star},\mathrm{C}^{\star},\mathrm{C^{\star}_{\alpha}}],$
(1)

where $T_{n}=(r_{n},x_{n})$ is composed of two components,
$r_{n}\in\mathrm{SO(3)}$ being a $3\times 3$ rotation matrix and
$x_{n}\in\mathbb{R}^{3}$ being a translation vector. Subsequently, all $N_{p}$
frames are denoted as
$\mathbf{T_{p}}=[T_{1},\dots,T_{N_{p}}]\in\mathrm{SE(3)}^{N_{p}}$. Lastly, one
may place each residue’s $\mathrm{O}^{\star}$ atom by rotating it around the
corresponding residue’s $C_{\alpha}-C$ bond using the torsion angle $\psi$.
For additional details regarding the mapping between idealized and absolute
protein atom positions, we refer readers to Yim et al. (2023).

Nucleic acid geometry parametrization. To model nucleic acid residues jointly
with protein residues, we require an expressive frame encoding of nucleic acid
geometry. Towards this end, we propose to model $N_{n}$ nucleic acid residues
as $N_{n}$ rigid body frames centered at each residue’s
$\mathrm{C4^{\prime}}^{\star}$ atom ($\mathrm{C4^{\prime}}^{\star}=(0,0,0)$),
with each frame representing a residue’s $\mathrm{C1^{\prime}}^{\star}$
through $\mathrm{C5^{\prime}}^{\star}$ atoms, $\mathrm{O3^{\prime}}^{\star}$
through $\mathrm{O5^{\prime}}^{\star}$ atoms, as well as its
$\mathrm{P}^{\star}$, $\mathrm{OP1}^{\star}$, $\mathrm{OP2}^{\star}$, and
$\mathrm{N9}^{\star}$ (or, for pyrimidine residues, $\mathrm{N1}^{\star}$)
atoms, respectively. Then for a given nucleic acid residue indexed by $n$, to
obtain the absolute positions of its main atoms, one may apply an
$\mathrm{SE(3)}$ transformation $T_{n}=(r_{n},x_{n})$ as

$[\mathrm{C3^{\prime}}_{n},\mathrm{O4^{\prime}}_{n},\mathrm{C4^{\prime}}_{n},\mathrm{C5^{\prime}}_{n}]=T_{n}\cdot[\mathrm{C3^{\prime}}^{\star},\mathrm{O4^{\prime}}^{\star},\mathrm{C4^{\prime}}^{\star},\mathrm{C5^{\prime}}^{\star}],$
(2)

where all $N_{n}$ frames are denoted as
$\mathbf{T_{n}}=[T_{N_{p}+1},\dots,T_{N_{n}}]\in\mathrm{SE(3)}^{N_{n}}$. For
nucleic acid residues, to place each residue’s auxiliary atoms (i.e.,
$\mathrm{C2^{\prime}}^{\star}_{n},\mathrm{C1^{\prime}}^{\star}_{n},\mathrm{N9}^{\star}_{n}(\mbox{or
}\mathrm{N1}^{\star}_{n}),\mathrm{O3^{\prime}}^{\star}_{n},\mathrm{O5^{\prime}}^{\star}_{n},\mathrm{P}^{\star}_{n},\mathrm{OP1}^{\star}_{n},\mbox{and
}\mathrm{OP2}^{\star}_{n})$, one may autoregressively rotate them around the
corresponding residue’s $\mathrm{C4^{\prime}}-\mathrm{C3^{\prime}}$,
$\mathrm{C4^{\prime}}-\mathrm{O4^{\prime}}$,
$\mathrm{O4^{\prime}}-\mathrm{C1^{\prime}}$,
$\mathrm{C4^{\prime}}-\mathrm{C3^{\prime}}$,
$\mathrm{C4^{\prime}}-\mathrm{C5^{\prime}}$,
$\mathrm{C5^{\prime}}-\mathrm{O5^{\prime}}$,
$\mathrm{O5^{\prime}}-\mathrm{P^{\prime}}$, and
$\mathrm{O5^{\prime}}-\mathrm{P^{\prime}}$ bonds using the torsion angles
$\Phi=[\phi_{1},\dots,\phi_{8}]$, respectively. For additional background
regarding nucleic acid geometry, we refer readers to Gelbin et al. (1996).

Joint protein-nucleic acid geometry parametrization. Given $N_{p}$ protein
frames and $N_{n}$ nucleic acid frames, respectively, we subsequently denote
the manifold of frames that we aim to represent with a generative model as
$\mathbf{N}=(N_{p},N_{n})$. Consequently, we recover the original means of
modeling a distribution $\mathbf{X}^{(0)}\sim p_{0}$ of frames in
$\mathrm{SE(3)}^{\mathbf{N}}$ supported on a Riemannian manifold
$\mathcal{M}$, as proposed by Yim et al. (2023). In short, using a score
network $s_{\theta}(t,\cdot)$, this allows one to model $\mathbf{X}^{(0)}$ via
an approximation of the Stein score $\nabla\log p_{t}$ using the
$\mathrm{SE(3)}$ denoising score matching (DSM) loss (Yim et al., 2023):

$\mathcal{L}(\theta)=\mathbb{E}[\lambda_{t}\lVert\nabla\log
p_{t|0}(\mathbf{X}^{(t)}|\mathbf{X}^{(0)})-s_{\theta}(t,\mathbf{X}^{(t)})\rVert^{2}],$
(3)

where $\theta$ are the network parameters to be optimized by minimizing
$\mathcal{L}$; $p_{t|0}$ is the density of $\mathbf{X}^{(t)}$ given
$\mathbf{X}^{(0)}$; $\lambda_{t}>0$ is a loss weight; and the expectation
$\mathbb{E}$ is taken over $(\mathbf{X}^{(0)},\mathbf{X}^{(t)})$ and
$t\sim\mathcal{U}([0,\mathrm{T_{F}}])$, with $\mathrm{T_{F}}$ being the final
timestep in a diffusion (i.e., noising) process over $\mathbf{X}^{(0)}$.

Remaining preliminaries. Following the notation of Yim et al. (2023), bold
denotes the concatenation of variables (e.g.,
$\mathbf{n}=(n_{1},\dots,n_{N})$); uppercase denotes stochastic variables
(e.g., $X\sim p$), whereas lowercase is reserved for deterministic variables;
and superscripts with parentheses denote time (e.g., $x^{(t)}$).

###  2.2 Joint Continuous-Discrete Diffusion in $\mathbb{R}^{3}$

Sequence modeling. Expressively modeling arbitrary macromolecular structures
is a non-trivial task for a generative model to address. Nonetheless, in this
work, we argue that an understanding of macromolecules is inherently
incomplete without also considering the sequence space of such molecules. In
particular, modeling of macromolecular sequences enables methods to learn
molecular representations that are complementary, and in many cases
foundational, to their structure knowledge of 3D molecular structures (Jumper
et al., 2021; Lin et al., 2022; Baek et al., 2022; Shen et al., 2022). This is
particularly relevant for nucleic acid modeling in that, practically speaking,
no general-purpose existing methods currently exist for fixed-backbone
sequence design of nucleic acid structures, although RNA-specific sequence
design methods are beginning to emerge Joshi et al. (2023). As such, in this
work, we are also interested in modeling the design space of macromolecular
sequences. We do so by introducing a discrete diffusion framework that, with
key adaptations to its noise scheduling, is amenable to the SE(3) Diffusion
framework of Yim et al. (2023) described above in Section 2.1. This allows our
proposed method, MMDiff, to jointly generate sequence-structure pairs of new
macromolecules, thereby circumventing the need for fixed-backbone sequence
design methods for nucleic acid residues given a generated backbone structure.

Discrete sequence generation. Diffusion generative modeling of discrete data
types such as amino acid and nucleic acid sequences has received considerably
less attention in the research literature compared to diffusion generative
approaches to modeling continuous data such as images or even 3D point cloud
coordinates. Nonetheless, preliminary works have investigated alternative
approaches to expressively modeling discrete values in a diffusion generative
setting (Hoogeboom et al., 2021; Austin et al., 2021; Chen et al., 2022), with
some directly adding discrete noise to one’s discrete network inputs and
others first representing discrete data in a continuous space and then
applying standard Gaussian diffusion techniques to generate new discrete data
samples through the use of an argmax operation (Ho et al., 2020). In this
work, we adopt the latter approach of modeling discrete macromolecular
sequence data as continuous representations such that (1) methodological
developments for continuous diffusion models will be directly amenable to such
an approach and (2) it then becomes trivial to condition sample generation on
arbitrary sequence features of interest (e.g., secondary structure annotations
for protein residues) by embedding such features in the same continuous space
as the input sequence. Concretely, inspired by Lisanza et al. (2023), we
propose to model a discrete amino acid and nucleic acid sequence
$\mathbf{S}_{d}$ in a shared vocabulary space (i.e., $\mathbb{R}^{s\times n}$)
as a zero-centered continuous one-hot vector representation

$\mathbf{S}_{c}^{(0)}=2*\mathrm{onehot}(\mathbf{S}_{d})-1=\\{2*(s\in\\{0.0,1.0\\}^{n}:\sum_{i=1}^{n}s_{i}=1.0)-1,\forall
s\in\mathbf{S}_{d})\\},$ (4)

such that discrete sequences can be generated by reversing a continuous
Gaussian diffusion process

$q(\mathbf{S}_{c}^{(t_{d})}|\mathbf{S}_{c}^{(0_{d})})=\mathcal{N}(\mathbf{S}_{c}^{(t_{d})}|\sqrt{\bar{\alpha}^{(t_{d})}}\mathbf{S}_{c}^{(0_{d})},(1-\bar{\alpha}^{(t_{d})})\mathrm{I})$
(5)

on an initially-random sequence vector representation
$\mathbf{S}_{c}^{(T_{d})}$ corresponding to the final sequence diffusion
timestep $T_{d}$. Here, $\bar{\alpha}^{(t_{d})}$ represents
$1-\beta^{(t_{d})}$, where $\beta^{(t_{d})}$ denotes one’s chosen variance
schedule for the sequence diffusion process, and $t_{d}$ represents a
discretized version of the structure diffusion continuous timestep $t$ via
equally-spaced binning of such timesteps. Note that this binning operation is
important to ensure that the sequence diffusion and structure diffusion
processes are in alignment with one another at each step in their respective
diffusion trajectories.

Training and sampling with structure noise. Reversal of the structure noise
schedule is performed by first using the FrameDiff algorithm of Yim et al.
(2023) to predict at each structure diffusion timestep the denoised frames
$\mathbf{F}^{(0)}$ to derive the respective rotation ($r$) and translation
($x$) scores
$\\{(\varsigma^{r}_{\theta,n},\varsigma^{x}_{\theta,n})\\}_{n=1}^{\mathbf{N}}$,
where during model training such score predictions are supervised using
Equation 3 along with supervision for predicted torsions $\Phi$ through
auxiliary backbone atom and inter-atom distance losses. The scores
$\\{(\varsigma^{r}_{\theta,n},\varsigma^{x}_{\theta,n})\\}_{n=1}^{\mathbf{N}}$
are then used to transition to a previous structure diffusion timestep as
$\mathbf{F}^{(t-1)}=\\{\exp_{\mathbf{F}_{n}^{(t)}}\\{(W_{n}^{r},W_{n}^{x})\\}\\}_{n=1}^{\mathbf{N}}$,
where $(W_{n}^{r},W_{n}^{x})$ are derived independently as
$W_{n}^{x}=P[P\gamma[\frac{1}{2}X_{n}^{(t)}+\varsigma_{\theta,n}^{x}]+\zeta\sqrt{\gamma}[Z_{n}^{x}\sim\mathcal{N}(0,\mathrm{Id}_{3})]]$
and
$W_{n}^{r}=\gamma\varsigma_{\theta,n}^{r}+\zeta\sqrt{\gamma}[Z_{n}^{r}\sim\mathcal{T}\mathcal{N}_{R_{n}^{(t)}}(0,\mathrm{Id})]$,
respectively (Yim et al., 2023). Here, $P\in\mathbb{R}^{3\mathbf{N}\times
3\mathbf{N}}$ represents a projection matrix that removes an object’s center
of mass; $\gamma=\frac{1-\epsilon}{N_{\mathrm{steps}}}$ with $\epsilon$ being
the final continuous timestep of the structure diffusion noise schedule prior
to predicting torsions $\Phi$; $X_{n}^{(t)}\in\mathbb{R}^{3}$ and
$R_{n}^{(t)}\in\mathbb{R}^{3\times 3}$ denote the translation vector and
rotation matrix associated with frame (i.e., residue) $\mathrm{F}_{n}^{(t)}$
at continuous timestep $t$; $Z_{n}^{x}$ and $Z_{n}^{r}$ denote Gaussian noise
sampled from the tangent spaces of $X_{n}^{(t)}$ and $R_{n}^{(t)}$,
respectively; and $\zeta\in[0,1]$ represents a scaling factor for structure
diffusion noise.

Training and sampling with sequence noise. As pointed out by Li et al. (2022);
Lisanza et al. (2023), in the context of training diffusion models for
discrete sequence generation it is also desirable to use a square root
diffusion noise schedule for $\beta^{(t_{d})}$, as using other diffusion noise
schedules during training may make it trivial for a model to predict a ground-
truth sequence from a small-noise discrete timestep $t_{d}$ using the argmax
operation. Such prior works thereby sensitize the model to sequence noise at
small $t_{d}$ timesteps by training it with a square root noise schedule,
while supervising its sequence predictions using a categorical cross-entropy
loss. After model training with a square root noise schedule, in this work we
also explore sampling sequences from a trained model using cosine and linear
noise schedules instead, to investigate how doing so affects the model at
inference time with regards to its sequence generation trajectories (and
thereby its structure generation trajectories), especially at small $t_{d}$
timesteps. For interested readers, these additional experiments are presented
in Appendix 5.2. Reversal of the chosen sequence noise schedule is then
performed by having a trained model directly predict
$\mathbf{S}_{c}^{(0_{d})}$ at timestep $t_{d}$ and then using this value in
Equation 5 to transition to timestep $t_{d}-1$.

###  2.3 Generating macromolecular complexes with MMDiff

As illustrated in Figure 1, our proposed method, MMDiff, enables joint
generation of nucleic acid sequences and structures through a hybrid
$\mathrm{SE(3)}$-discrete diffusion process. Moreover, by way of its design as
described in Section 2.1, MMDiff can simultaneously model protein sequences
and structures together with nucleic acids, enabling the generation of inter-
molecular complexes. In the remainder of this section, we will discuss the
remaining details necessary to train and sample from MMDiff for generating
macromolecular complexes.

Model architecture, data, and featurization. Inspired by its demonstrated
ability to generate designable, diverse, and novel monomeric protein
structures, for MMDiff we adapt the FrameDiff model architecture of Yim et al.
(2023) for modeling of macromolecular complexes. To train MMDiff for
macromolecular modeling, we downloaded from the RCSB PDB (Berman et al., 2000)
all available instances of protein-nucleic acid structures using the
PDBManager API within the Graphein Python package (Jamasb et al., 2022). We
filtered the downloaded structures to those that (1) were determined at a
resolution of 4.5 Å or better using nuclear magnetic resonance, diffraction,
or cryo-electron microscopy; (2) contain protein chains between lengths 40 and
256 or nucleic acid chains between lengths 3 and 256; and (3) contain no more
than 10 protein or 10 nucleic acid chains. In Appendix 5.3, we plot the length
and chain composition distributions of the various training dataset
configurations we explore in this work. To accurately characterize the
intricate chain compositions present within many macromolecular PDB complexes,
we subsequently adopt the relative position encoding of AlphaFold-Multimer
(Evans et al., 2021) in contrast to the absolute position encoding used by Yim
et al. (2023). Otherwise, we adapt the feature initialization scheme of Trippe
et al. (2022) to accommodate both proteins and nucleic acids.

Intra-molecule consensus sampling. To generate internally-consistent
macromolecular sequences, we propose a consensus sampling algorithm that,
during sampling, constrains MMDiff to select only amino acid residue types in
a designated protein sequence region and only the residues types of a single
type of nucleic acid molecule (i.e., either DNA or RNA) in a nucleic acid
sequence region. Namely, for a given sequence representation
$\mathbf{S}_{c}^{(t_{d})}$, we propose to consensus-sample sequences by
dynamically modifying the pseudo-probabilities (i.e., continuous
representation) of $\mathbf{S}_{c}^{(t_{d})}$ such that the largest
probabilities associated with residue types outside of the current sequence’s
molecule type consensus (e.g.,
$C^{(t_{d})}=[\textsc{DA},\textsc{DC},\textsc{DG},\textsc{DT}]=$ DNA residue
types if 50% or more of the generated residues are DNA residues) are set to be
marginally smaller than the smallest probabilities associated with residue
types within the current molecule type consensus as

$\mathbf{S}_{c}^{(t_{d})}=\\{s\in\mathbb{R}^{n}:\max_{i\notin
C^{(t_{d})}}s_{i}<\min_{i\in C^{(t_{d})}}s_{i},\ \forall
s\in\mathbf{S}_{c}\mbox{ with }i\in[1,n]\\}.$ (6)

Applying Equation 6 after each iteration of sequence denoising during sampling
effectively biases the generative sequence distribution from which MMDiff
samples new sequences, such that the model is increasingly less likely to
naturally sample out-of-vocabulary within a given chain.

##  3 Experiments

Method | $<5\text{\r{A}}\ \mathrm{scRMSD}$ | $\mathrm{D}_{chain}$ | $\mathrm{D}_{complex}$ | $\mathrm{D}_{single}$ | $\mathrm{D}_{all}$ | $\mathrm{N}_{pool}$
---|---|---|---|---|---|---
Random Generation | 0.00% | 0.92 | 1.00 | 1.00 | 1.00 | 0.85
MMDiff | 0.74% | 0.92 | 1.00 | 1.00 | 1.00 | 0.72
MMDiff-Protein | 0.00% | 0.91 | 0.99 | 0.99 | 0.99 | 0.73
(a) Protein Generation

Method | $<5\text{\r{A}}\ \mathrm{scRMSD}$ | $\mathrm{D}_{chain}$ | $\mathrm{D}_{complex}$ | $\mathrm{D}_{single}$ | $\mathrm{D}_{all}$ | $\mathrm{N}_{pool}$
---|---|---|---|---|---|---
Random Generation | 1.33% | 0.99 | 1.00 | 0.99 | 0.99 | 0.75
MMDiff | 8.67% | 1.00 | 1.00 | 1.00 | 1.00 | 0.74
MMDiff-NA | 6.00% | 0.98 | 0.98 | 0.99 | 0.99 | 0.65
(b) Nucleic Acid Generation

Method | $<5\text{\r{A}}\ \mathrm{scRMSD}$ | $\mathrm{D}_{chain}$ | $\mathrm{D}_{complex}$ | $\mathrm{D}_{single}$ | $\mathrm{D}_{all}$ | $\mathrm{N}_{pool}$
---|---|---|---|---|---|---
Random Generation | 0.37% | 1.00 | 1.00 | 1.00 | 1.00 | 0.72
MMDiff | 0.74% | 0.96 | 0.85 | 1.00 | 1.00 | 0.69
MMDiff-Monomer | 0.00% | 0.96 | 0.85 | 1.00 | 1.00 | 0.70
(c) Protein-Nucleic Acid Complex Generation

Table 1: A comparison of our proposed method, MMDiff, against a random
baseline for generating different types of biological complexes, where we
assess each method’s designability (based on an $\mathrm{scRMSD}$ threshold),
diversity (regarding four variants of $\mathrm{D}$), and novelty (i.e.,
$\mathrm{N}_{pool}$). Across all tasks, MMDiff represents the same model
trained on protein-nucleic acid complexes. However, in task (a), MMDiff-
Protein is trained on both protein monomers and protein complexes; in task
(b), MMDiff-NA is trained solely on nucleic acid complexes; and lastly, in
task (c), MMDiff-Monomer is trained on both monomers and protein-nucleic acid
complexes.

Metrics. To investigate the performance of MMDiff (and similar generative
models developed in the future) in generating new macromolecular complexes, we
propose to examine the quality of its generated complexes from multiple
perspectives: (1) the model’s ability to generate designable sequences such
that, upon predicting the structure of such a designed sequence using an
external structure prediction model, the difference between the model’s co-
designed structure and the sequence-predicted structure is minimal; (2) the
model’s ability to generate diverse complex structures that are distinct from
each other after structure-based clustering; and (3) the model’s ability to
generate (3) novel complex structures that are well-differentiated from
complex structures found within the model’s training dataset.

Designability. To instantiate such metrics in the context of macromolecular
complex generation as outlined above, for (1) we adopt RoseTTAFold2NA (Baek et
al., 2022) for sequence-based prediction of generated protein-nucleic acid
structures to assess designability, where we report structural similarities in
terms of self-consistency RMSD and TM-score (i.e., $\mathrm{scRMSD}$ and
$\mathrm{scTM}$) for each complex as determined by US-align (Zhang et al.,
2022). Due to practical concerns regarding the significant computational
requirements of running RoseTTAFold2NA for each complex, in this setting, we
generate and evaluate one sequence per structure produced by the model, using
RoseTTAFold2NA in single-sequence mode to reduce its average prediction time
from 25 minutes to 2 minutes. Note that this may limit the sequence
designability performance that MMDiff can achieve, as previous works have
found it necessary to design up to 8 sequences per generated structure to
achieve good designability results. Moreover, RoseTTAFold2NA was trained to
predict protein-nucleic acid structures using multiple sequence alignments
(MSAs) as a crucial input feature, hence running the model in single-sequence
mode will incur some prediction accuracy degradation in exchange for practical
running times. Taking all of these constraints into account, we propose to
evaluate each method’s ability to generate molecules with an $\mathrm{scRMSD}$
$<5$ Å to denote a successfully-designed macromolecule, with an
$\mathrm{scRMSD}$ $<2$ Å being a more ideal threshold (Yim et al., 2023) yet
being largely impractical due to the potential mismatches between generated
and predicted inter-molecule interaction types (e.g., RF2NA predicting non-
interacting chain structures and MMDiff generating interacting chain
structures).

Figure 2: Comparison of $\mathrm{scRMSD}$ complex designability results using
different training methods. Here, the top row corresponds to protein-only
experiments, the middle row to nucleic acid-only experiments, and the bottom
row to protein-nucleic acid experiments. The columns denotes samples generated
using the random macromolecule generation baseline, MMDiff, and
MMDiff-{Protein, NA, Monomer} (corresponding to the {first, second, third}
row), respectively. Note that novel data samples are displayed with a *
symbol. Overall, most designable complexes contain 2-3 chains, and most
generated complexes contain novel chains (with a novelty $>0.7$).

Diversity. Regarding (2), we use qTMclust (Zhang et al., 2022) for
macromolecular structure-based clustering of the model’s pool of generated
complex structures to measure diversity, where the diversity of a pool is
defined in four components as follows. We define the baseline $\mathrm{chain\
diversity}$ ($\mathrm{D}_{chain}$) of a pool as: (number of chain clusters) /
(number of chains). Similarly, as an upper bound on the model’s diversity, we
define the simpler $\mathrm{complex\ diversity}$ ($\mathrm{D}_{complex}$) of a
pool as: (number of complex clusters) / (number of complexes). More
specifically, to measure how often the model generates a unique interaction
type, we then define the $\mathrm{single-chain\ complex\ diversity}$
($\mathrm{D}_{single}$) as: (number of complexes for which one of its chains
is a chain cluster representative) / (number of complexes). Lastly, as a lower
bound on the model’s diversity, we define the $\mathrm{all-chain\ complex\
diversity}$ ($\mathrm{D}_{all}$) as: (number of complexes for which each of
its chains is a chain cluster representative) / (number of complexes).

Novelty. To address (3) we perform pairwise structural alignments between the
model’s pool of generated complex structures and the model’s training dataset
using US-align (Zhang et al., 2022) to estimate novelty, where the novelty of
a pool is quantified as the inverse of the maximum TM-score between a
generated complex structure and any complex structure in the training dataset
(i.e., $\mathrm{N}_{pool}=1-\mathrm{maxTM}$).

Baselines. To the best of our knowledge, no prior methods can co-design
sequences and structures of macromolecular complexes. As such, in addition to
evaluating MMDiff, we report results using a naive baseline that randomly
generates sequences and structures of protein, nucleic acid, and protein-
nucleic acid complexes, respectively, as a simple baseline to evaluate a
macromolecular generative model’s performance. For the random baseline, we
extract the initial sequence and structure noise representations of MMDiff’s
diffusion process (i.e., noise representations at timestep $T$) as the random
baseline’s generated outputs. As such, one should expect a generative model
that has learned meaningful sequence and structural patterns from its training
dataset to outperform such a baseline regarding the designability of its
samples by fully reversing its learned diffusion process.

(a) Designed micro-RNA.

(b) Designed single-stranded DNA.

(c) Designed protein-DNA.

Figure 3: Examples of macromolecules successfully designed by MMDiff.

Results and Analysis. As shown in Table 1(a) and Figure 2, the results for
protein complex generation suggest that MMDiff (trained on protein-nucleic
acid complexes) can successfully generate a handful of designable protein
complexes, whereas the random baseline cannot. Interestingly, MMDiff-Protein,
a version of MMDiff that is trained not only on protein-nucleic acid complexes
but also jointly on the large monomeric protein structure dataset of Yim et
al. (2023), underperforms MMDiff in all design metrics. One explanation for
this observation is that the generative modeling space that MMDiff-Protein is
tasked to model when jointly training on protein monomers is much larger and
noisier compared to the modeling space of solely protein-nucleic acid
complexes, a notion that is supported by the striking dataset distribution
shift observed between Figures 7 and 10.

Table 1(b) shows that both MMDiff (trained on protein-nucleic acid complexes)
and MMDiff-NA (trained solely on nucleic acid complexes) outperform the random
baseline method for nucleic acid complex generation in terms of designability
(i.e., $\mathrm{scRMSD}$ and $\mathrm{scTM}$). Important to note is that
MMDiff generalizes markedly beyond the designability of the random baseline by
generating a reasonable number of moderately designable, highly diverse, and
novel complex structures. Figure 2 illustrates the designability scores MMDiff
can achieve in terms of $\mathrm{scRMSD}$ across a variety of complex lengths
and chain compositions. Lastly, Figures 3(a) and 3(b) illustrate micro-RNA
(GCGCGCGGGG) and single-stranded micro-DNA (TCTTGGTTTGTCTTTTCGA) molecules
successfully designed using MMDiff, which achieve $\mathrm{scRMSD}$ measures
of 1.26 and 4.95 and maximum trainTM measures of 0.36 and 0.31 with respect to
the PDB complexes 4WTF and 1XPX, respectively. Moreover, RoseTTAFold2NA is
confident in its predicted structure for these designed sequences, yielding
per-residue average plDDT scores of 95.11 and 78.24, respectively. Note, for
illustrating these molecules, the corresponding MMDiff-generated structures
are shown in green (a) or orange (b), whereas the RF2NA-predicted structures
are shown in orange (a) or green (b).

Table 1(c) and Figure 2 show that MMDiff outperforms the random baseline as
well as MMDiff-Monomer, similar to the results listed in Table 1(a) for
protein generation. We argue this result is also likely due to increased
modeling space complexity introduced by adding a majority of protein monomers
to the training dataset of MMDiff-Monomer. Figure 3(c) illustrates a protein-
DNA complex (DNA sequence: GGGGGG) designed by MMDiff. Note that, here, MMDiff
generates a structurally-novel protein-DNA interaction in comparison to that
of RF2NA while maintaining the main structure of the DNA molecule
($\mathrm{scRMSD}$ of DNA structure $<5$ Å), where RF2NA’s per-residue average
plDDT of 58.10 may indicate the possibility of multiple conformations for such
an interaction type (as suggested by MMDiff). Note, for visualizing this
macromolecule, the MMDiff-generated chains are shown in green (orange),
whereas the RF2NA-predicted chains are shown in pink (purple). Also, note that
all corresponding designability results in terms of $\mathrm{scTM}$ are
displayed in Appendix 5.1 within Figure 4.

###  3.1 Related Work

Generative methods for molecule generation. Several kinds of generative models
have been proposed for molecular structure generation (Hoogeboom et al., 2022;
Eguchi et al., 2022; Xu et al., 2023) and sequence generation (Hsu et al.,
2022), most commonly diffusion generative models (Anand and Achim, 2022; Wu et
al., 2022; Trippe et al., 2022; Lin and AlQuraishi, 2023; Fu et al., 2023;
Watson et al., 2023). Regarding diffusion generative models, most closely
related to this work is the $\mathrm{SE(3)}$ diffusion model of Yim et al.
(2023) and the joint protein sequence-structure diffusion model of Lisanza et
al. (2023). The former is designed solely to generate protein backbone
structures, and the latter method uses a pre-trained structure prediction
network to design protein sequence-structure pairs using pure Gaussian noise.
In contrast to these methods, MMDiff designs sequence-structure pairs of not
only proteins but also nucleic acid complexes using a geometrically-expressive
rigid body denoising process for structure generation and a discrete diffusion
process for sequence generation, thereby deriving the benefits of both data
representations for macromolecular inputs while avoiding the need for pre-
training.

Nucleic acid modeling. Few works have rigorously explored geometric modeling
of nucleic acid structures to the extent that methods such as AlphaFold 2
(Jumper et al., 2021) have done for amino acid residues. Some preliminary
works in this direction are those of Shen et al. (2022); Baek et al. (2022);
Joshi et al. (2023). Nonetheless, no prior works have presented a geometric
modeling scheme for nucleic acids that is amenable to macromolecular design
tasks while being complementary to amino acid modeling, which we present in
this work.

##  4 Discussions & Conclusions

In this work, we introduced MMDiff, a macromolecular sequence-structure
diffusion generative model for nucleic acid and protein complexes. Through
rigorous experiments within our proposed computational benchmark for
macromolecular design, we have demonstrated that MMDiff can generate
designable, diverse, and novel nucleic acid structures.

Limitations. Currently, MMDiff achieves many more success cases for nucleic
acid design compared to protein(-nucleic acid) design, suggesting room for
improvement with regards to joint structural modeling of both types of
molecules. Our preliminary investigations into this topic suggest that
acquiring new sources of diverse and high-quality macromolecular training data
will be key for future developments, as the PDB is currently quite limited in
the number of protein-nucleic acid quaternary structures it offers. We argue
this based on our observation that extended training of MMDiff on protein-
nucleic acid complexes, over time, leads the model to increase its
designability metrics for nucleic acids yet decrease its designability metrics
for proteins and protein-nucleic acid complexes, hinting at a phenomenon akin
to overfitting.

Future work. Given the limitations above, future work could involve (1)
curating large new training datasets containing protein complexes with or
without DNA and RNA molecules for joint or separate generation tasks; (2)
extending MMDiff to support full-atom protein-nucleic design in an end-to-end
manner; (3) validating MMDiff’s nucleic acid designs through wet-lab
experiments; or (4) investigating efficient ways of scoring generative models
producing large multi-chain molecular assemblies such as protein-RNA and
protein-DNA complexes. In this work, we have already observed the difficulties
of assessing such macromolecular design methods, so we believe standardized
benchmarking amongst future methods will be critical to ensure progress in the
field.

## Acknowledgments and Disclosure of Funding

The authors would like to thank Jason Yim for insightful discussions regarding
the $\mathrm{SE(3)}$ Diffusion framework and its successful applications in
protein backbone design. This work was funded by Profluent Bio and performed
while Alex Morehead was a summer research intern with the company.

## References

  * Anand and Achim (2022) Namrata Anand and Tudor Achim.  Protein structure and sequence generation with equivariant denoising diffusion probabilistic models.  _arXiv preprint arXiv:2205.15019_ , 2022.
  * Austin et al. (2021) Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg.  Structured denoising diffusion models in discrete state-spaces.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, _Advances in Neural Information Processing Systems_ , volume 34, pages 17981–17993. Curran Associates, Inc., 2021.  URL https://proceedings.neurips.cc/paper_files/paper/2021/file/958c530554f78bcd8e97125b70e6973d-Paper.pdf.
  * Baek et al. (2022) Minkyung Baek, Ryan McHugh, Ivan Anishchenko, David Baker, and Frank DiMaio.  Accurate prediction of nucleic acid and protein-nucleic acid complexes using rosettafoldna.  _bioRxiv_ , pages 2022–09, 2022.
  * Berman et al. (2000) 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.
  * Chen et al. (2022) Ting Chen, Ruixiang Zhang, and Geoffrey Hinton.  Analog bits: Generating discrete data using diffusion models with self-conditioning.  _arXiv preprint arXiv:2208.04202_ , 2022.
  * Eguchi et al. (2022) Raphael R Eguchi, Christian A Choe, and Po-Ssu Huang.  Ig-vae: Generative modeling of protein structure by direct 3d coordinate generation.  _PLoS computational biology_ , 18(6):e1010271, 2022.
  * 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_ , pages 2021–10, 2021.
  * Fu et al. (2023) Cong Fu, Keqiang Yan, Limei Wang, Wing Yee Au, Michael McThrow, Tao Komikado, Koji Maruhashi, Kanji Uchino, Xiaoning Qian, and Shuiwang Ji.  A latent diffusion model for protein structure generation.  _arXiv preprint arXiv:2305.04120_ , 2023.
  * Gelbin et al. (1996) Anke Gelbin, Bohdan Schneider, Lester Clowney, Shu-Hsin Hsieh, Wilma K Olson, and Helen M Berman.  Geometric parameters in nucleic acids: sugar and phosphate constituents.  _Journal of the American Chemical Society_ , 118(3):519–529, 1996.
  * Hie and Yang (2022) Brian L. Hie and Kevin K. Yang.  Adaptive machine learning for protein engineering.  _Current Opinion in Structural Biology_ , 72:145–152, 2022.  ISSN 0959-440X.  doi: https://doi.org/10.1016/j.sbi.2021.11.002.  URL https://www.sciencedirect.com/science/article/pii/S0959440X21001457.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in neural information processing systems_ , 33:6840–6851, 2020.
  * Hoogeboom et al. (2021) Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling.  Argmax flows and multinomial diffusion: Learning categorical distributions.  _Advances in Neural Information Processing Systems_ , 34:12454–12465, 2021.
  * Hoogeboom et al. (2022) Emiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling.  Equivariant diffusion for molecule generation in 3d.  In _International conference on machine learning_ , pages 8867–8887. PMLR, 2022.
  * Hsu et al. (2022) Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie, Tom Sercu, Adam Lerer, and Alexander Rives.  Learning inverse folding from millions of predicted structures.  In _International Conference on Machine Learning_ , pages 8946–8970. PMLR, 2022.
  * Jamasb et al. (2022) Arian Jamasb, Ramon Viñas Torné, Eric Ma, Yuanqi Du, Charles Harris, Kexin Huang, Dominic Hall, Pietro Lió, and Tom Blundell.  Graphein-a python library for geometric deep learning and network analysis on biomolecular structures and interaction networks.  _Advances in Neural Information Processing Systems_ , 35:27153–27167, 2022.
  * Joshi et al. (2023) Chaitanya K. Joshi, Arian R. Jamasb, Ramon Viñas, Charles Harris, Simon Mathis, and Pietro Liò.  Multi-state rna design with geometric multi-graph neural networks.  _arXiv preprint arXiv:2305.14749_ , 2023.
  * Jumper et al. (2021) John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al.  Highly accurate protein structure prediction with alphafold.  _Nature_ , 596(7873):583–589, 2021.
  * Li et al. (2022) Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto.  Diffusion-lm improves controllable text generation.  _Advances in Neural Information Processing Systems_ , 35:4328–4343, 2022.
  * Lin and AlQuraishi (2023) Yeqing Lin and Mohammed AlQuraishi.  Generating novel, designable, and diverse protein structures by equivariantly diffusing oriented residue clouds.  _arXiv preprint arXiv:2301.12485_ , 2023.
  * Lin et al. (2022) Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Sal Candido, et al.  Language models of protein sequences at the scale of evolution enable accurate structure prediction.  _BioRxiv_ , 2022:500902, 2022.
  * Lisanza et al. (2023) Sidney Lyayuga Lisanza, Jacob Merle Gershon, Sam Wayne Kenmore Tipps, Lucas Arnoldt, Samuel Hendel, Jeremiah Nelson Sims, Xinting Li, and David Baker.  Joint generation of protein sequence and structure with rosettafold sequence space diffusion.  _bioRxiv_ , pages 2023–05, 2023.
  * Lutz (2010) Stefan Lutz.  Beyond directed evolution—semi-rational protein engineering and design.  _Current Opinion in Biotechnology_ , 21(6):734–743, 2010.  ISSN 0958-1669.  doi: https://doi.org/10.1016/j.copbio.2010.08.011.  URL https://www.sciencedirect.com/science/article/pii/S0958166910001540.  Chemical biotechnology – Pharmaceutical biotechnology.
  * Shen et al. (2022) Tao Shen, Zhihang Hu, Zhangzhi Peng, Jiayang Chen, Peng Xiong, Liang Hong, Liangzhen Zheng, Yixuan Wang, Irwin King, Sheng Wang, et al.  E2efold-3d: end-to-end deep learning method for accurate de novo rna 3d structure prediction.  _arXiv preprint arXiv:2207.01586_ , 2022.
  * Trippe et al. (2022) Brian L Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi Jaakkola.  Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem.  _arXiv preprint arXiv:2206.04119_ , 2022.
  * Watson et al. (2023) Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al.  De novo design of protein structure and function with rfdiffusion.  _Nature_ , pages 1–3, 2023.
  * Wu et al. (2022) Kevin Eric Wu, Kevin K Yang, Rianne van den Berg, James Zou, Alex Xijie Lu, and Ava P Amini.  Protein structure generation via folding diffusion.  _arXiv preprint arXiv:2209.15611_ , 2022.
  * Xu et al. (2023) Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec.  Geometric latent diffusion models for 3d molecule generation.  In _International Conference on Machine Learning_ , pages 38592–38610. PMLR, 2023.
  * 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 et al. (2022) Chengxin Zhang, Morgan Shine, Anna Marie Pyle, and Yang Zhang.  Us-align: universal structure alignments of proteins, nucleic acids, and macromolecular complexes.  _Nature methods_ , 19(9):1109–1115, 2022.

##  5 Supplementary Material

###  5.1 Additional Designability Results

Figure 4: Comparison of $\mathrm{scTM}$ complex designability results using
different training methods. Here, the top row corresponds to protein-only
experiments, the middle row to nucleic acid-only experiments, and the bottom
row to protein-nucleic acid experiments. The columns denotes samples generated
using the random macromolecule generation baseline, MMDiff, and
MMDiff-{Protein, NA, Monomer} (corresponding to the {first, second, third}
row), respectively. Note that novel data samples are displayed with a *
symbol. Overall, most designable complexes contain 2-3 chains, and most
generated complexes contain novel chains (with a novelty $>0.7$).

In Figure 4, we report designability results in terms of $\mathrm{scTM}$ for
different training methods as a complement to the results in Figure 2.

###  5.2 Alternative Noise Schedules for Sequence Generation

(a) $\mathrm{scRMSD}$ designability results for nucleic acid structures using
a version of MMDiff trained only on nucleic acid complexes and sampled using
linear sequence noise.

(b) $\mathrm{scRMSD}$ designability results for nucleic acid structures using
a version of MMDiff trained only on nucleic acid complexes and sampled using
cosine sequence noise.

Figure 5: Comparison of $\mathrm{scRMSD}$ complex designability results for
nucleic acid structures using different sequence noise schedules during
sampling. Zoom in for the best viewing experience.

(a) $\mathrm{scTM}$ designability results for nucleic acid structures using a
version of MMDiff trained only on nucleic acid complexes and sampled using
linear sequence noise.

(b) $\mathrm{scTM}$ designability results for nucleic acid structures using a
version of MMDiff trained only on nucleic acid complexes and sampled using
cosine sequence noise.

Figure 6: Comparison of $\mathrm{scTM}$ complex designability results for
nucleic acid structures using different sequence noise schedules during
sampling. Zoom in for the best viewing experience.

As Figures 5 and 6 illustrate, in terms of designability metrics, linear and
cosine sequence noise schedules marginally outperform a square root noise
schedule, producing slightly more designable macromolecules overall. However,
to simplify the presentation of results, in the main text we report MMDiff
generation results using a square root noise schedule.

###  5.3 Dataset Distributions

Figure 7: The distribution of protein-nucleic acid complex lengths upon which
MMDiff is trained. Figure 8: The distribution of protein monomer and protein
complex lengths upon which MMDiff is trained. Figure 9: The distribution of
nucleic acid complex lengths upon which MMDiff is trained. Figure 10: The
distribution of protein-nucleic acid complex (and monomer) lengths upon which
MMDiff-Monomer (for protein-nucleic acid complex generation) is trained.

(a) The radius of gyration training distribution for nucleic acid structures.

(b) The radius of gyration generated distribution for nucleic acid structures.

Figure 11: A comparison between the radius of gyration distributions of
MMDiff’s training dataset and MMDiff’s generated samples.

Figure 7 displays the length distribution of protein-nucleic acid complexes
within MMDiff’s protein-nucleic acid training dataset, whereas Figures 8 and 9
display (in isolation) the length distributions of protein complexes and
nucleic acid complexes, respectively, within the protein-nucleic acid training
dataset. Likewise, Figure 10 illustrates the length distribution of protein-
nucleic acid complexes (and protein monomers) within MMDiff’s training dataset
after combining it with the monomeric protein structure dataset of Yim et al.
[2023].

Lastly, to investigate how well MMDiff can model the true radius of gyration
distribution for nucleic acid structures in its training dataset, Figures
11(a) and 11(b) display the distributions of (1) MMDiff’s training dataset and
(2) MMDiff’s generated samples. As shown in these figures, MMDiff generates
its samples with an average radius of gyration of 1.8, which is reasonably
close to the 1.65 average radius of gyration of its training dataset.


# Knowledge Graph Error Detection with Contrastive Confidence Adaption

Xiangyu Liu1, Yang Liu1, Wei Hu1,2, Corresponding author

###### Abstract

Knowledge graphs (KGs) often contain various errors. Previous works on
detecting errors in KGs mainly rely on triplet embedding from graph structure.
We conduct an empirical study and find that these works struggle to
discriminate noise from semantically-similar correct triplets. In this paper,
we propose a KG error detection model CCA to integrate both textual and graph
structural information from triplet reconstruction for better distinguishing
semantics. We design interactive contrastive learning to capture the
differences between textual and structural patterns. Furthermore, we construct
realistic datasets with semantically-similar noise and adversarial noise.
Experimental results demonstrate that CCA outperforms state-of-the-art
baselines, especially in detecting semantically-similar noise and adversarial
noise.

## Introduction

A knowledge graph (KG) is composed of triplets in the form of $(head\
entity,relation,tail\ entity)$, which finds extensive applications in
downstream tasks like question answering (Saxena, Tripathi, and Talukdar 2020)
and recommender systems (Guo et al. 2022). Existing KGs such as NELL (Carlson
et al. 2010) and Knowledge Vault (Dong et al. 2014) continuously extract
triplets in an automatic way, which inevitably introduces noise. Detecting
these errors holds the potential to improve the quality of KGs.

Existing works on KG error detection can be classified into embedding-based
and path-based models. The former (Bordes et al. 2013; Yang et al. 2015;
Trouillon et al. 2016) learns confidence scores based on the representations
of entities and relations. The latter (Lin et al. 2015; Jia et al. 2019) uses
paths between entities to evaluate the confidence of triplets. Different from
the task of link prediction (Chen et al. 2021) or triplet classification (Yao,
Mao, and Luo 2019), error detection focuses on detecting the error triplets in
the whole unsupervised KG, aiming to capture the variance of triplets and give
an accurate estimate of their confidence.

Current KG error detection models face a significant challenge due to the
unavailability of noise patterns and the difficulty in acquiring accurately
labeled noise samples for robust supervision. Negative sampling by replacing
entities is widely used in previous works, especially the embedding-based
models. However, real-world scenarios often introduce confusing noise
semantically related to correct samples. Let us see two real error triplets in
the FB15K-237 dataset (Toutanova et al. 2015): (George Lopez, profession, Disc
jockey) and (Majel Barrett, profession, Writer). In the former, the correct
tail entity should be ‘Comedian’, and in the latter, the correct tail entity
should be ‘Actress’. Notably, ‘Disc jockey’ and ‘Writer’ both represent
professions, mirroring common human errors. This form of noise is harder to
differentiate and aligns closely with human error tendencies.

Models |  Noise types | FB15K-237 | WN18RR
---|---|---|---
$K$=1% |  $K$=5% |  $K$=1% |  $K$=5%
CAGED | Random | 0.945 | 0.758 | 0.795 | 0.486
Similar | 0.633 | 0.367 | 0.657 | 0.384
TransE | Random | 0.904 | 0.726 | 0.630 | 0.434
Similar | 0.611 | 0.303 | 0.503 | 0.361
Table 1: Results of our empirical study. Error triplets are divided into the
“random” and “similar” groups, based on the methods of replacing entities. We
show the top-$K$ precision of two typical models on FB15K-237 and WN18RR.
Figure 1: An overview of the proposed model CCA. (a) BERT and Transformer-
based graph encoders extract textual and graph structural information,
respectively. (b) The reconstruction module classifies error triplets by
reconstructing head and tail entities in textual and structural embedding. (c)
Interactive contrastive learning aligns the projection of textual and
structural embeddings and recognizes errors by inter-model difference. (d) The
knowledge fusion module takes pseudo labels generated from aggregated results
as triplet confidence, which is further injected into the training process.

We conduct a further empirical study to explore the performance of existing
works on two specific types of noise: random noise and semantically-similar
noise. Random noise is generated by randomly replacing the head or tail entity
of a correct triplet, which is used in previous works (Dong et al. 2023).
Semantically-similar noise uses entities that have co-occurrence with their
relations, which means that it is semantically related to the relation. We
test two typical models CAGED (Zhang et al. 2022) and TransE (Bordes et al.
2013) on the FB15K-237 (Toutanova et al. 2015) and WN18RR (Dettmers et al.
2018) datasets to verify whether existing works can deal with more realistic
noise. As shown in Table 1, we add 5% of the two types of noise to the
datasets and show the precision of the top-1% and top-5% detected triplets.
The results show that, although the existing methods CAGED and TransE perform
well on random noise, their effects are greatly reduced on semantically-
similar noise. This is because existing models predominantly rely on graph
structure, ignoring the rich textual information of KGs. The incompleteness of
KGs leads to the lack of important information that can distinguish
semantically-similar noise, which exhibits greater semantic relevance and
shares a similar graph structure with correct triplets.

Furthermore, we consider the prevalent scenarios of KG error detection to
build adversarial noise. For automatic KG construction and completion, noise
is inevitably introduced. An effective error detection model should possess
the capability to pinpoint triplets that have been inaccurately completed or
constructed. We filter out the error results from the construction model to
constitute adversarial noise. Subsequent experiments also show that textual
information plays a key role in the noise generated by completion which relies
on graph structure.

Existing works need a potent method to discern more realistic noise and
extract the full potential of textual information within KGs for error
detection. To achieve this goal, we opt to leverage a pre-trained language
model (PLM) as the encoder of textual information in our model. We refer to
KG-BERT (Yao, Mao, and Luo 2019), PKGC (Lv et al. 2022), MLMLM (Clouatre et
al. 2021), and CoLE (Liu et al. 2022), which perform well in KG embedding and
link prediction. A PLM uses a large number of open-domain corpora for training
and can supplement rich information for triplets.

We propose a novel KG error detection model CCA. It leverages the
reconstruction of triplets to comprehend noise patterns from both textual and
graph structural perspectives. Also, we design interactive contrastive
learning to align the latent representations of textual and structural
information. It facilitates noise identification based on disparities between
these two forms of information. CCA combines the reconstruction and
contrastive learning output and generates pseudo-labels to represent triplet
confidence. This adaptive confidence guides model training by alleviating
noise interference and transferring knowledge between reconstruction and
contrastive learning. With the utilization of textual information, CCA not
only outperforms the state-of-the-art methods on random noise but also
performs more prominently in semantically-similar noise and adversarial noise,
validating the effectiveness of CCA in complex real-world scenarios.

In summary, this paper makes the following contributions:

  * •

We propose an end-to-end KG error detection model, which fully leverages both
textual and structural information by reconstructing triplets, and alleviates
the interference of noise. It transfers the knowledge between reconstruction
and contrastive learning.

  * •

We design interactive contrastive learning to align the latent representations
of textual and structural information. We use different negative sampling
policies to mine anomalous features on the alignment of latent spaces.

  * •

We construct two kinds of noise, semantically-similar noise and adversarial
noise, to evaluate the performance of our model in more realistic scenarios.
Experiments show that CCA not only surpasses the state-of-the-art competitors
on random noise but also achieves more significant results on the datasets
with semantically-similar noise and adversarial noise. Datasets and source
code are available at https://github.com/nju-websoft/CCA.

## Related Work

### KG Error Detection

The embedding-based models, e.g., TransE (Bordes et al. 2013), DistMult (Yang
et al. 2015), and ComplEx (Trouillon et al. 2016), evaluate triplets by
designing score functions according to the representations of entities and
relations. TransE employs $\lVert\mathbf{h}+\mathbf{r}-\mathbf{t}\rVert$ as
the score function to assess triplets. Subsequent works such as CKRL (Xie et
al. 2018), CAGED (Zhang et al. 2022), and SCEF (Zhao and Liu 2019) adopt the
TransE’s score function as a component. The embedding-based models heavily
rely on negative sampling, and their effectiveness is hampered by the
challenge of accurately modeling the real noise distribution. Consequently,
the negative samples often deviate from the real noise, undermining the error
detection performance.

Some works exploit path information in KGs as a reference. PTransE (Lin et al.
2015) integrates relation path information into triplet embedding learning.
KGTtm (Jia et al. 2019) leverages path information for error detection.

Several works recognize that relying solely on graph structure may not be
sufficient for robust KG error detection. Integrating additional information
can offer more semantic insights into triplets or provide more accurate
supervision signals, thereby enhancing model capabilities. Defacto (Lehmann et
al. 2012) uses the information from relevant web pages to verify triplets,
while the categorical information of entities and relations is also useful to
assess the trustworthiness of triplets (Paulheim and Bizer 2014). CrossVal
(Wang, Ma, and Gao 2020) introduces human-curated knowledge repositories to
assess triplet confidence in the target KG through the correlation of triplets
in two KGs. Crowdsourcing and active learning models, such as KAEL (Dong et
al. 2023), show promising results. CKRL (Xie et al. 2018) and PGE (Cheng et
al. 2022) also explore confidence-aware methods. Their confidence constraints
are based solely on the score function of the TransE’s loss.

The existing models mainly focus on utilizing graph structure while under-
utilizing valuable textual information. By considering this, our proposed
model CCA seeks to overcome these limitations by jointly leveraging textual
and graph structural information, aiming to achieve more accurate and
comprehensive KG error detection.

### PLMs for KGs

While KG error detection using PLMs is relatively scarce, the use of PLMs has
demonstrated success in various tasks like KG completion. KG-BERT (Yao, Mao,
and Luo 2019) and PKGC (Lv et al. 2022) achieve good results by directly fine-
tuning PLMs by the triplet classification tasks using carefully constructed
input templates. MMLLM (Clouatre et al. 2021) effectively employs the [MASK]
token in prompts to enhance link prediction.

To fuse language models with structural models, (Nadkarni et al. 2021; Wang et
al. 2021; Chen et al. 2023) leverage the text enhancement method to represent
entities by combining textual and structural features. CoLE (Liu et al. 2022)
explores the use of the two types of information to construct models
independently and uses knowledge distillation to complement and improve each
other’s performance.

## The Proposed Model

In this section, we describe the proposed model CCA in detail. The framework
is shown in Figure 1. The goal of the proposed model is to better leverage
both textual and graph structural information at the same time. Given a
triplet $(h,r,t)$, we first construct its input sequence by its subgraph and
the descriptions of the entities and the relation. Then, BERT and a
Transformer-based encoder are used to extract textual features and graph
structural features, respectively. Reconstruction loss is employed to
construct the head or tail entity in the triplet and evaluate the triplet’s
trustworthiness. Interactive contrastive learning utilizes the projection of
textual and structural representations to discover anomalous features in the
two alignment spaces. The scores from contrastive learning and reconstruction
classifier are aggregated to generate a pseudo label as the confidence to
dynamically constrain the training process.

### Feature Extraction

Text encoder. Currently PLMs like BERT (Kenton and Toutanova 2019) have been
widely used to extract textual features from KGs. For each entity $e$,
following the previous work (Liu et al. 2022), we leverage its description
$D_{e}$ and human-readable literal name $N_{e}$ to learn the textual
representation. We first add the corresponding representation of the new
entity into the PLM’s vocabulary. To better leverage the description of the
entity, pre-trained prompts are constructed for initializing the new token
embedding $\mathbf{e}$:

$\displaystyle\mathbf{e}=\mathrm{BERT}(\text{``The description of [MASK] is
$D_{e}$''}),$ (1)

where [MASK] denotes the missing entity and we use its embedding to initialize
the token embedding $\mathbf{e}$. After pre-training on the description, we
construct two prompts through the masked head and tail entities to better
leverage textual information.

Given a triplet $(h,r,t)$, we use the prompts presented in Table 2 as the
input sequence. In $P_{h}^{T}$, the [MASK] token is used to replace the
original head entity, and the tail entity $t$ in $P_{t}^{T}$ is also masked,
so that $h$ and $t$ can be reconstructed respectively by the encoder. [CLS]
and [SEP] are two special tokens used to separate the different parts of the
prompt. [SEP]${}_{i}^{r}$ is the $i$-th adjustable soft prompt token for the
relation $r$, as a more expressive separator. We provide the textual
description $D_{t}$ of the tail entity in $P_{h}^{T}$, and $D_{h}$ of the head
entity is used in $P_{t}^{T}$. If the tail entity $t$ is an anomalous entity,
it is hard to reconstruct $t$ through the text description of $h$ in
$P_{t}^{T}$, so the reconstruction loss can be used to evaluate the error
possibility of triplets from the textual view.

$P_{t}^{T}$ = [CLS] [SEP]${}_{1}^{r}$ $N_{h}$ [SEP]${}_{2}^{r}$ $N_{r}$
[SEP]${}_{3}^{r}$ [MASK] [SEP]${}_{4}^{r}$ $D_{h}$
---
$P_{h}^{T}$ = [CLS] [SEP]${}_{1}^{r}$ [MASK] [SEP]${}_{2}^{r}$ $N_{r}$
[SEP]${}_{3}^{r}$ $N_{t}$ [SEP]${}_{4}^{r}$ $D_{t}$
$P_{t}^{S}$ = [CLS] [$h$] [SEP] [$r$] [SEP] [MASK] [SEP] [$h^{\prime}$]
[$r^{\prime}$]…
$P_{h}^{S}$ = [CLS] [MASK] [SEP] [$r$] [SEP] [$t$] [SEP] [$t^{\prime}$]
[$r^{\prime}$]…
Table 2: Inputs used in the textual and structural encoders

Structure encoder. To make full use of the graph structure in KGs, we build an
entity’s subgraph by random sampling and use Transformer to encode it. Given a
triplet $(h,r,t)$, we also create two input sequences $P_{h}^{S}$ and
$P_{t}^{S}$ by masking $h$ and $t$, respectively. For $P_{h}^{S}$, let
$InNeighbor(h)=\big{\\{}(h^{\prime},r^{\prime})\,|\,(h^{\prime},r^{\prime},h)\in\Gamma\big{\\}}$
and
$OutNeighbor(h)=\big{\\{}(t^{\prime},r^{\prime})\,|\,(h,r^{\prime},t^{\prime})\in\Gamma\big{\\}}$
denote two neighbor sets of entity $h$, where $\Gamma$ is the triplet set. We
randomly select some neighbors and separate them by the [SEP] token, as shown
in Table 2. $P_{t}^{S}$ is constructed in the same way.

The structure encoder can learn the differences between $t$ and the subgraph
of $h$. Therefore, by reconstructing entities from the subgraphs, errors can
be distinguished.

### Reconstruction Classifier

Given a triplet $(h,r,t)$, we can get the output representation
$\mathbf{e}_{h}^{T}$ of the masked head entity from the text encoder:

$\displaystyle\mathbf{e}_{h}^{T}=\mathrm{BERT}(P_{h}^{T}),$ (2)

where $P_{h}^{T}$ is the input sequence that we construct in Table 2.

Then, we can calculate the prediction logits against all candidate entities
and get the cross-entropy loss $\mathcal{L}_{h}^{T}$ according to the
corresponding labels:

$\displaystyle\mathcal{L}_{h}^{T}=CE\Big{(}\mathrm{softmax}\big{(}\mathrm{MLP}(\mathbf{e}_{h}^{T}),\mathbf{E}^{T}\big{)},\mathbf{y}_{h}\Big{)},$
(3)

where $\mathbf{E}^{T}$ is the matrix of all entity embeddings and
$\mathbf{y}_{h}$ is the corresponding labels. $\mathbf{e}_{h}^{T}$ is
transferred by a multi-layer perceptron and used to calculate the prediction
logits with all embeddings in $\mathbf{E}^{T}$. $CE()$ denotes the cross-
entropy loss. Note that the tail entity should also be reconstructed, and we
calculate $\mathbf{e}_{t}^{T}$ and $\mathcal{L}_{t}^{T}$ in the same way.

Next, we introduce adaptive confidence to the reconstruction loss and obtain
the final text reconstruction loss:

$\displaystyle\mathcal{L}_{\mathrm{text}}=\sum_{(h,r,t)\in\Gamma}c(h,r,t)\big{(}\mathcal{L}_{h}^{T}+\mathcal{L}_{t}^{T}\big{)},$
(4)

where the triplet confidence $c(h,r,t)$ is calculated in Section Knowledge
Fusion shortly. Note that we use the same way on the Transformer encoder to
get the masked entity losses $\mathcal{L}_{h}^{S}$ and $\mathcal{L}_{t}^{S}$,
and the overall reconstruction loss $\mathcal{L}_{struct}$ similar to Eqs.
(2), (3), and (4). We add $\mathcal{L}_{text}$ and $\mathcal{L}_{struct}$ and
jointly train the final reconstruction loss
$\mathcal{L}_{\mathrm{reconstruct}}$ as follows:

$\displaystyle\mathcal{L}_{\mathrm{reconstruct}}=\alpha\mathcal{L}_{\mathrm{text}}+(1-\alpha)\mathcal{L}_{\mathrm{struct}},$
(5)

where $\alpha$ is a hyperparameter to balance the training process.

Finally, we use the cross-entropy loss to compute the scores of the triplet
$(h,r,t)$:

$\displaystyle score_{\mathrm{text}}$
$\displaystyle=\mathcal{L}_{h}^{T}+\mathcal{L}_{t}^{T},$ (6) $\displaystyle
score_{\mathrm{struct}}$
$\displaystyle=\mathcal{L}_{h}^{S}+\mathcal{L}_{t}^{S},$ (7)

where $score_{\mathrm{text}}$ and $score_{\mathrm{struct}}$ are the confidence
scores from text and structure reconstruction, respectively.

### Interactive Contrastive Learning

We use interactive contrastive learning (Yang et al. 2022) to learn the
differences between textual and structural information. For a triplet
$(h,r,t)$, $\mathbf{v}_{i}^{h}$ and $\mathbf{v}_{i}^{t}$ represent the token
embeddings generated from the PLM by $P_{h}^{T}$ and $P_{t}^{T}$,
respectively. $\mathbf{u}_{i}^{h}$ and $\mathbf{u}_{i}^{t}$ represent the
embeddings from the structure encoder. $\mathbf{v}_{i}^{h}$ contains the
textual information of $t$ and the reconstruction of $h$ from the PLM.
$\mathbf{u}_{i}^{t}$ has the structural information of $h$ and the
reconstruction of $t$ from Transformer. Intuitively, $\mathbf{v}_{i}^{h}$ and
$\mathbf{u}_{i}^{t}$ should be aligned, so that the prediction of $h$ can
match its structural information and the prediction of $t$ can match its
textual information. Therefore, we use $\mathbf{v}_{i}^{h}$ and
$\mathbf{u}_{i}^{t}$ as one anchor pair and $\mathbf{v}_{i}^{t}$ and
$\mathbf{u}_{i}^{h}$ as the other pair for contrastive learning.

Negative sampling. In interactive contrastive learning, we maximize the
agreement between the anchor pair like $\mathbf{v}_{i}^{h}$ from the textual
encoder and its structural representation $\mathbf{u}_{i}^{t}$ for one triplet
$(h,r,t)$. We use two negative sampling policies to support stable training.

First, we construct negative samples to align the latent spaces from the two
encoders. For the anchor sample $\mathbf{v}_{i}^{h}$ from the text encoder,
another sample in the structure encoder like $\mathbf{u}_{j}^{t}$ should keep
a distance from the anchor sample. Thus, we randomly take another sample from
the structure encoder as one of the negative samples.

To improve the sensitivity and robustness of the model against noise, we
construct error samples as negative examples based on the anchor. We randomly
replace the head (or tail) entity of the anchor sample $(h,r,t)$ to construct
a typical error sample $(h,r,t^{\prime})$ (or $(h^{\prime},r,t)$), and
generate its representations $\mathbf{m}_{i_{k}}^{h}$ and
$\mathbf{n}_{i_{k}}^{t}$ from different encoders, respectively. Note that
$(h,r,t^{\prime})$ is in fact the noise generated from disturbance. Therefore,
$\mathbf{m}_{i_{k}}^{h}$ and $\mathbf{n}_{i_{k}}^{t}$ should not be aligned.

To reduce the training cost, we directly use the representation of the entity
obtained by the encoder in each training batch to disturb original embeddings.
For a triplet $(h,r,t)$, the similarity of its negative samples is calculated
as

$\displaystyle\begin{split}neg\\_sim(h,r,t)&=\sum\limits_{j\neq
i}\exp\big{(}sim(\mathbf{v}_{i}^{h},\mathbf{u}_{j}^{t})\big{)}\\\
&\quad+\sum\limits_{x=1}^{X}\exp\big{(}sim(\mathbf{m}_{i_{x}}^{h},\mathbf{n}_{i_{x}}^{t})\big{)},\end{split}$
(8)

where $\mathbf{v}_{i}^{h}$ denotes the masked head triplet embedding of the
$i$-th sample and $\mathbf{u}_{j}^{t}$ denotes the masked tail triplet
embedding of the $j$-th sample. $sim()$ denotes the cosine similarity used to
calculate the distance of the two embeddings. $\mathbf{m}_{i_{x}}^{h}$ and
$\mathbf{n}_{i_{x}}^{t}$ represent the randomly replaced embeddings from
$\mathbf{v}_{i}^{h}$ and $\mathbf{u}_{i}^{t}$, and we replace them for $X=4$
times randomly.

Adaptive contrastive learning. We employ the InfoNCE loss (van den Oord, Li,
and Vinyals 2018) to train interactive contrastive learning. For the anchor
pair $\mathbf{v}_{i}^{h}$ and $\mathbf{u}_{i}^{t}$, the loss
$\mathcal{L}_{\mathrm{ICL}}^{1}(h,r,t)$ is

(9)

where $neg\\_sim(h,r,t)$ is the similarity of negative samples. Note that if
we choose $\mathbf{v}_{i}^{t}$ and $\mathbf{u}_{i}^{h}$ as the anchor pair, we
can get $\mathcal{L}_{\mathrm{ICL}}^{2}(h,r,t)$ in the same way. Therefore,
the final contrastive training loss is

(10)

where $c(h,r,t)$ is the adaptive confidence of $(h,r,t)$. The score function
of contrastive learning of $(h,r,t)$ is

$\displaystyle
score_{\mathrm{contrastive}}=sim(\mathbf{v}_{i}^{h},\mathbf{u}_{j}^{t}).$ (11)

### Knowledge Fusion

Since $score_{\mathrm{text}}$ and $score_{\mathrm{struct}}$ are both logits
from triplet reconstruction, we directly add them to obtain the reconstruction
score:

$\displaystyle
score_{\mathrm{reconstruct}}=score_{\mathrm{text}}+\lambda\,score_{\mathrm{struct}},$
(12)

where $\lambda$ is the hyperparameter to balance the scores. Note that the
scores from reconstruction and contrastive learning have quite different
distributions, so we use ranking to combine them. By ranking the error scores
$score_{\mathrm{reconstruct}}$ and $score_{\mathrm{contrastive}}$, we obtain
the error ranks $R_{1}(h,r,t)$ and $R_{2}(h,r,t)$ of each triplet,
respectively. The final score is

$\displaystyle
score(h,r,t)=\frac{1}{\lceil\frac{R_{1}(h,r,t)}{\gamma}\rceil}+\frac{1}{\lceil\frac{R_{2}(h,r,t)}{\gamma}\rceil^{\beta}},$
(13)

where $\gamma$ is the hyperparameter to control the number of triplets with
the same score. $\beta$ is used to balance the scores.

We generate a pseudo label from the final score. Let
$Z=normalize\big{(}[z_{1},\dots,z_{n}]\big{)}$ be the pseudo label set, where
$z_{i}\sim N(\mu,\rho)$ and $Z$ is sorted in ascending order, the adaptive
confidence of $(h,r,t)$ is as follows:

$\displaystyle c(h,r,t)=z_{R(h,r,t)},$ (14)

where $R(h,r,t)$ is the rank of $score(h,r,t)$ and we use $c(h,r,t)$ as the
final result. Note that $c(h,r,t)$ is the result integrated from textual and
structural reconstruction and contrastive learning. We use $c(h,r,t)$ directly
constrain the training process according to Eqs. (4) and (10) to transfer
knowledge across components.

## Experiments and Results

### Dataset Construction

We conduct our experiments on FB15K-237 (Toutanova et al. 2015) and WN18RR
(Dettmers et al. 2018). We add 5% noise into these two datasets. The
statistics of the new datasets are shown in Table 3. We construct three types
of noise to better evaluate the performance of our model:

Datasets | Correct triplets | Wrong triplets | Avg. deg.
---|---|---|---
FB15K-237 | 310,116 | 16,321 | 44.89
WN18RR |  93,003 |  4,894 |  4.78
Table 3: Statistics of the modified datasets

  * •

Random noise, drawn from a uniform distribution that does not depend on the
data and is not predictable. For a correct triplet $(h,r,t)$, random noise is
constructed by randomly replacing one of the entities or the relation.

  * •

Semantically-similar noise, which is more realistic because errors primarily
stem from semantic confusion in real-world scenarios. To better evaluate
semantically-related error detection, we introduce semantically-similar noise.
Given a triplet $(h,r,t)$, we form a candidate set $S_{t}$ by selecting other
tail entities linked to $r$. For each entity $e\in S_{t}$, we employ BERT to
encode its description, acquiring the semantic embedding $\mathbf{e}$ of $e$.
We calculate the sampling probability distribution of semantic similarity as
follows and utilize it as a sampling probability to replace the original $t$:

$\displaystyle\mathbf{P}(t)=\mathrm{softmax}\Big{(}\big{[}\mathbf{t}\cdot\mathbf{e}_{1},\mathbf{t}\cdot\mathbf{e}_{2},\dots,\mathbf{t}\cdot\mathbf{e}_{i}\big{]}\Big{)},$
(15)

where $\mathbf{t}$ denotes the embedding of the original tail entity $t$,
$\mathbf{e}_{i}$ denotes the embedding of the $i$-th candidate entity, and
$\cdot$ denotes the dot product.

  * •

Adversarial noise, which is adversarially generated from KG construction
models. Since error detection models are frequently employed to discern errors
in automatically constructed KGs, we use TransE (Bordes et al. 2013) for
adversarial noise generation and evaluate the ability of models to recognize
errors during KG construction. Given a dataset $D$, we randomly divide it into
training and testing sets, $D_{train}$ and $D_{test}$. After training on
$D_{train}$, we randomly select one entity within the top-10 prediction
results of the triplet in $D_{test}$ to construct its error triplet. We
iteratively repeat this process until we get an adequate amount of noise.

|  Model types | Models | FB15K-237 | WN18RR
---|---|---|---|---
|  $K$=1% |  $K$=2% |  $K$=3% |  $K$=4% |  $K$=5% |  $K$=1% |  $K$=2% |  $K$=3% |  $K$=4% |  $K$=5%
Precision@top-$K$ |  Struct models | TransE | 0.946 | 0.774 | 0.606 | 0.498 | 0.423 | 0.690 | 0.576 | 0.501 | 0.437 | 0.400
DistMult | 0.764 | 0.630 | 0.530 | 0.463 | 0.410 | 0.687 | 0.633 | 0.526 | 0.438 | 0.374
ComplEx | 0.802 | 0.633 | 0.521 | 0.446 | 0.393 | 0.774 | 0.699 | 0.550 | 0.449 | 0.384
KGTtm | 0.857 | 0.687 | 0.631 | 0.467 | 0.437 | 0.789 | 0.644 | 0.541 | 0.473 | 0.417
CAGED | 0.863 | 0.666 | 0.602 | 0.543 | 0.467 | 0.753 | 0.620 | 0.536 | 0.470 | 0.421
Text models | KG-BERT | 0.966 | 0.799 | 0.660 | 0.584 | 0.498 | 0.973 | 0.968 | 0.938 | 0.829 | 0.710
StAR | 0.970 | 0.835 | 0.681 | 0.571 | 0.490 | 0.971 | 0.918 | 0.842 | 0.739 | 0.647
CSProm-KG | 0.961 | 0.798 | 0.689 | 0.574 | 0.509 | 0.977 | 0.927 | 0.869 | 0.773 | 0.680
CCA (ours) | 0.969 | 0.812 | 0.707 | 0.599 | 0.534 | 0.986 | 0.959 | 0.920 | 0.834 | 0.733
Recall@top-$K$ |  Struct models | TransE | 0.189 | 0.310 | 0.364 | 0.399 | 0.423 | 0.138 | 0.231 | 0.300 | 0.350 | 0.400
DistMult | 0.153 | 0.252 | 0.318 | 0.371 | 0.410 | 0.137 | 0.253 | 0.316 | 0.350 | 0.374
ComplEx | 0.161 | 0.254 | 0.313 | 0.357 | 0.393 | 0.155 | 0.279 | 0.330 | 0.359 | 0.384
KGTtm | 0.171 | 0.275 | 0.378 | 0.374 | 0.437 | 0.158 | 0.257 | 0.324 | 0.378 | 0.417
CAGED | 0.173 | 0.266 | 0.362 | 0.435 | 0.467 | 0.150 | 0.248 | 0.321 | 0.376 | 0.421
Text models | KG-BERT | 0.193 | 0.319 | 0.396 | 0.467 | 0.498 | 0.195 | 0.387 | 0.563 | 0.663 | 0.710
StAR | 0.194 | 0.334 | 0.409 | 0.457 | 0.490 | 0.194 | 0.367 | 0.505 | 0.591 | 0.647
CSProm-KG | 0.192 | 0.319 | 0.413 | 0.459 | 0.509 | 0.195 | 0.371 | 0.521 | 0.618 | 0.680
CCA (ours) | 0.194 | 0.325 | 0.424 | 0.479 | 0.534 | 0.197 | 0.384 | 0.552 | 0.667 | 0.733
Table 4: Results of precision and recall at top-$K$ on FB15K-237 and WN18RR

### Settings

We use the ranking measures for evaluation. All triplets in a KG are ranked
based on the confidence scores in ascending order. Triplets with lower
confidence scores are more likely to be noisy. We follow CAGED (Zhang et al.
2022) and use precision@top-$K$ and recall@top-$K$ to assess the performance,
where $K$ denotes the ratio (e.g., 5%).

All experiments are conducted on two Intel Xeon Gold 6326 CPUs, 512GB RAM, and
one NVIDIA RTX A6000 GPU. We leverage the BERT-base model from huggingface as
the PLM. We use PyTorch to implement our model and employ the AdamW optimizer
and a cosine decay scheduler with a linear warm-up for optimization. The grid
search is used for hyperparameter tuning. The results of KG embedding models
are obtained from $\mu$KG (Luo, Sun, and Hu 2022), a recent open-source
library for KG embedding.

### Baseline Models

We compare our model with eight baselines, including five structural models
and three textual models.

  * •

Structural models. We choose three typical embedding models TransE (Bordes et
al. 2013), DistMult (Yang et al. 2015), and ComplEx (Trouillon et al. 2016)
and two state-of-the-art structural error detection models KGTtm (Jia et al.
2019) and CAGED (Zhang et al. 2022) for comparison. Both KGTtm and CAGED
utilize graph structural information based on the TransE’s score function.

  * •

Textual models. We choose one textual classification model KG-BERT (Yao, Mao,
and Luo 2019), and two recent models StAR (Wang et al. 2021) and CSProm-KG
(Chen et al. 2023), which leverage both textual and structural information,
for comparison. KG-BERT takes entity and relation descriptions of a triplet as
input and computes scores by binary classification. StAR and CSProm-KG use the
representations obtained from PLMs to aid structural models. CCA also falls
into this category.

| Models |  $K$=1% |  $K$=3% |  $K$=5%
---|---|---|---|---
FB15K-237 | CCA (full) | .969 / .194 | .707 / .424 | .535 / .535
– adapt conf. | .951 / .190 | .664 / .398 | .496 / .496
– inter contr. | .961 / .192 | .679 / .407 | .509 / .509
– struct recon. | .797 / .159 | .582 / .349 | .475 / .475
– text recon. | .778 / .156 | .543 / .325 | .414 / .414
WN18RR | CCA (full) | .986 / .197 | .920 / .552 | .733 / .733
– adapt conf. | .971 / .194 | .858 / .515 | .632 / .632
– inter contr. | .979 / .196 | .911 / .546 | .722 / .722
– struct recon. | .974 / .195 | .914 / .549 | .726 / .726
– text recon. | .577 / .115 | .513 / .308 | .423 / .423
Table 5: Ablation results of precision and recall at top-$K$

### Main Results

Table 4 presents the comparison results of our model and eight baseline models
on FB15K-237 and WN18RR, where we add 5% noise, containing three types of
noise in equal quantities. Overall, our model outperforms eight baseline
models on both datasets. We have three observations below:

First, compared with the KG embedding models, the error detection models
generally perform better. This is because the KG embedding models assume that
all triplets are correct. They learn representations of entities and relations
without alleviating disturbance from noise, making it difficult to
discriminate error triplets. With adaptive confidence, CCA can effectively
improve performance.

Second, benefiting from PLMs, the textual models outperform the structural
models on the two datasets. They both can learn noise patterns in the textual
view with the descriptions of entities and relations, considering that PLMs
can capture factual knowledge from a large amount of open-domain corpora. CCA
combines both textual information and structural information and uses a
Transformer encoder to learn structural information from scratch, which is
more effective than other textual models.

| Models | Random | Similar | Adversarial
---|---|---|---|---
FB15K-237 | TransE | 0.726 | 0.304 | 0.125
ComplEx | 0.734 | 0.306 | 0.150
DistMult | 0.662 | 0.328 | 0.125
KGTtm | 0.730 | 0.318 | 0.128
CAGED | 0.758 | 0.331 | 0.126
KG-BERT | 0.674 | 0.336 | 0.199
StAR | 0.728 | 0.371 | 0.164
CSProm-KG | 0.732 | 0.419 | 0.211
CCA (ours) | 0.768 | 0.453 | 0.240
WN18RR | TransE | 0.434 | 0.351 | 0.314
ComplEx | 0.384 | 0.303 | 0.366
DistMult | 0.316 | 0.338 | 0.285
KGTtm | 0.448 | 0.391 | 0.359
CAGED | 0.486 | 0.388 | 0.373
KG-BERT | 0.806 | 0.633 | 0.599
StAR | 0.794 | 0.610 | 0.527
CSProm-KG | 0.791 | 0.634 | 0.536
CCA (ours) | 0.807 | 0.657 | 0.545
Table 6: Precision at top-5% on three different error types |  |  | CCA | KG-BERT | CAGED
---|---|---|---|---|---
Case 1. | Noise | Razzie Award for Worst Actor, award winner, Richard D. Zanuck | 2.9% |  7.6% |  3.4%
Truth | Razzie Award for Worst Actor, award winner, Marlon Wayans
Case 2. | Noise |  Tony Award for Best Choreography, ceremony, 54th Academy Awards | 3.2% | 10.6% | 11.6%
Truth | Academy Award for Best Sound Mixing, ceremony, 54th Academy Awards
Case 3. | Noise |  Lauren Tom, place of birth, Buenos Aires | 2.2% |  3.2% | 35.1%
Truth | Sebastian Krys, place of birth, Buenos Aires
Case 4. | Noise | Viola Davis, award nominee, Carrie-Anne Moss | 1.0% | 10.9% |  4.3%
Truth | Viola Davis, award nominee, Meryl Streep
Table 7: Case study on FB15K-237

Third, the performance gap between CCA and KG-BERT on WN18RR is less
significant than that on FB15K-237. We think the reason is that WN18RR is much
sparser than FB15K-237. As shown in Table 3, the average degree of entities in
WN18RR is far less than that of FB15K-237, which indicates that there is less
structural information in WN18RR. Thus, our model obtains less improvement by
combining with graph structure on sparse KGs.

### Ablation Study

We conduct an ablation study to assess the impact of each component in CCA. We
have four variants of our model by removing confidence adaption, interactive
contrastive learning, structure reconstruction, or text reconstruction. We
report their performance under the same settings.

Table 5 shows that all components contribute to our model on both datasets,
where text reconstruction contributes the most. On FB15K-237, the precision
improvement of text construction is from 0.414 to 0.535. This shows that
applying PLMs to KG error detection is a promising way, and it deserves
further research. The improvement is more obvious on WN18RR. The reason is
that the structure-based models perform poorly, as WN18RR is sparser than
FB15K-237.

Removing adaptive confidence affects differently on the two datasets. On
WN18RR, adaptive confidence contributes more than FB15K-237, as the
performance gap between text and structure reconstruction is larger. Without
adaptive confidence, structure reconstruction is hard to gain benefits from
the knowledge of PLMs, which provides scores with lower accuracy and disturbs
the final results. Therefore, adaptive confidence can be more effective in
improving overall performance when components have a larger performance gap.

### Performance on Different Error Types

To investigate the robustness of all models to different noise types, we add
three types of 5% noise separately to the datasets. Note that precision@top-5%
and recall@top-5% have the same value as we add 5% noise. Table 6 presents
precision@top-5% on FB15K-237 and WN18RR.

On FB15K-237, CCA outperforms all baselines. For random noise, the structure-
based models generally outperform KG-BERT, as this type of noise can be well
distinguished by graph structure alone. For semantically-similar noise, all
baseline models perform closely. As for adversarial noise, the textual models
outperform the structural ones, because this type of noise consists of
entities that are wrongly predicted as correct entities by TransE, which is a
structural model. Our CCA takes advantage of the knowledge from PLMs and graph
structure at the same time, so it outperforms all baselines, especially for
semantically-similar noise.

On WN18RR, CCA is comparable to KG-BERT, which is different from the
observations on FB15K-237. The models with PLMs largely outperform the
structural models on all three types of noise, mainly due to that WN18RR is
sparser than FB15K-237. For adversarial noise, our model underperforms KG-
BERT. Given that the structure-based models do not work well on WN18RR,
Transformer in CCA provides less reliable knowledge for the PLM, hindering it
from distinguishing the noise correctly.

### Case Study

To explore how textual information and graph structure act on error detection,
we perform a case study on CCA, KG-BERT, and CAGED. Table 7 shows the position
of the error triplet in a ranking, where a smaller percentage is better.

In the first two cases, CCA leverages both textual and graph structural
information, and outperforms KG-BERT and CAGED, which solely use one type of
information. For example, in Case 1, the description shows that Richard D.
Zanuck is an American film producer, and the graph structural information
records films that he has produced. Evidence from text and graph structure are
complementary to each other. In Case 3, CAGED is not effective compared with
CCA and KG-BERT, which is caused by the lack of graph structural information.
The degrees of entities “Lauren Tom” and “The Venture Bros” are 15 and 17,
respectively, which are much smaller than the average degree of 44.89 in
FB15K-237. In Case 4, entities in the noise and correct triplets all have
abundant graph structural information, so CAGED can achieve a better effect.

## Conclusion

In this paper, we propose a novel KG error detection model. It encodes textual
and graph structural information to find noise patterns. To alleviate the
disturbance of noise and integrate the knowledge from the two encoders, we
design a confidence adaption model to aggregate the results and constrain the
training process. To learn the noise patterns between textual and structural
information, we leverage interactive contrastive learning to align latent
spaces. We construct semantically-similar noise and adversarial noise for
evaluation. Experiments show that our model achieves good results on
semantically-similar noise and adversarial noise.

## Acknowledgments

This work is supported by the National Natural Science Foundation of China
(No. 62272219).

## References

  * Bordes et al. (2013) Bordes, A.; Usunier, N.; García-Durán, A.; Weston, J.; and Yakhnenko, O. 2013.  Translating embeddings for modeling multi-relational data.  In _NIPS_ , 2787–2795.
  * Carlson et al. (2010) Carlson, A.; Betteridge, J.; Kisiel, B.; Settles, B.; Hruschka, E.; and Mitchell, T. 2010.  Toward an architecture for never-ending language learning.  In _AAAI_ , 1306–1313.
  * Chen et al. (2023) Chen, C.; Wang, Y.; Sun, A.; Li, B.; and Lam, K.-Y. 2023.  Dipping PLMs sauce: Bridging structure and text for effective knowledge graph completion via conditional soft prompting.  In _Findings of ACL_ , 11489–11503.
  * Chen et al. (2021) Chen, S.; Liu, X.; Gao, J.; Jiao, J.; Zhang, R.; and Ji, Y. 2021.  HittER: Hierarchical transformers for knowledge graph embeddings.  In _EMNLP_ , 10395–10407.
  * Cheng et al. (2022) Cheng, K.; Li, X.; Xu, Y. E.; Dong, X. L.; and Sun, Y. 2022.  PGE: Robust product graph embedding learning for error detection.  _Proc. VLDB Endow._ , 15(6): 1288–1296.
  * Clouatre et al. (2021) Clouatre, L.; Trempe, P.; Zouaq, A.; and Chandar, S. 2021.  MLMLM: Link prediction with mean likelihood masked language model.  In _Findings of ACL_ , 4321–4331.
  * Dettmers et al. (2018) Dettmers, T.; Minervini, P.; Stenetorp, P.; and Riedel, S. 2018.  Convolutional 2D knowledge graph embeddings.  In _AAAI_ , 1811–1818.
  * Dong et al. (2023) Dong, J.; Zhang, Q.; Huang, X.; Tan, Q.; Zha, D.; and Zihao, Z. 2023.  Active ensemble learning for knowledge graph error detection.  In _WSDM_ , 877–885.
  * Dong et al. (2014) Dong, X.; Gabrilovich, E.; Heitz, G.; Horn, W.; Lao, N.; Murphy, K.; Strohmann, T.; Sun, S.; and Zhang, W. 2014.  Knowledge vault: A web-scale approach to probabilistic knowledge fusion.  In _KDD_ , 601–610.
  * Guo et al. (2022) Guo, Q.; Zhuang, F.; Qin, C.; Zhu, H.; Xie, X.; Xiong, H.; and He, Q. 2022.  A survey on knowledge graph-based recommender systems.  _IEEE Trans. Knowl. Data Eng._ , 34(8): 3549–3568.
  * Jia et al. (2019) Jia, S.; Xiang, Y.; Chen, X.; and Wang, K. 2019.  Triple trustworthiness measurement for knowledge graph.  In _WWW_ , 2865–2871.
  * Kenton and Toutanova (2019) Kenton, J. D. M.-W. C.; and Toutanova, L. K. 2019.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _NAACL_ , 4171–4186.
  * Lehmann et al. (2012) Lehmann, J.; Gerber, D.; Morsey, M.; and Ngomo, A.-C. N. 2012.  Defacto-deep fact validation.  In _ISWC_ , 312–327.
  * Lin et al. (2015) Lin, Y.; Liu, Z.; Luan, H.; Sun, M.; Rao, S.; and Liu, S. 2015.  Modeling relation paths for representation learning of knowledge bases.  In _EMNLP_ , 705–714.
  * Liu et al. (2022) Liu, Y.; Sun, Z.; Li, G.; and Hu, W. 2022.  I know what you do not know: Knowledge graph embedding via co-distillation learning.  In _CIKM_ , 1329–1338.
  * Luo, Sun, and Hu (2022) Luo, X.; Sun, Z.; and Hu, W. 2022.  $\mu$KG: A library for multi-source knowledge graph embeddings and applications.  In _ISWC_ , 610–627.
  * Lv et al. (2022) Lv, X.; Lin, Y.; Cao, Y.; Hou, L.; Li, J.; Liu, Z.; Li, P.; and Zhou, J. 2022.  Do pre-trained models benefit knowledge graph completion? A reliable evaluation and a reasonable approach.  In _Findings of ACL_ , 3570–3581.
  * Nadkarni et al. (2021) Nadkarni, R.; Wadden, D.; Beltagy, I.; Smith, N.; Hajishirzi, H.; and Hope, T. 2021\.  Scientific language models for biomedical knowledge base completion: An empirical study.  In _AKBC_.
  * Paulheim and Bizer (2014) Paulheim, H.; and Bizer, C. 2014.  Improving the quality of linked data using statistical distributions.  _Int. J. Semant. Web Inf. Syst._ , 10(2): 63–86.
  * Saxena, Tripathi, and Talukdar (2020) Saxena, A.; Tripathi, A.; and Talukdar, P. P. 2020.  Improving multi-hop question answering over knowledge graphs using knowledge base embeddings.  In _ACL_ , 4498–4507.
  * Toutanova et al. (2015) Toutanova, K.; Chen, D.; Pantel, P.; Poon, H.; Choudhury, P.; and Gamon, M. 2015\.  Representing text for joint embedding of text and knowledge bases.  In _EMNLP_ , 1499–1509.
  * Trouillon et al. (2016) Trouillon, T.; Welbl, J.; Riedel, S.; Gaussier, É.; and Bouchard, G. 2016.  Complex embeddings for simple link prediction.  In _ICML_ , 2071–2080.
  * van den Oord, Li, and Vinyals (2018) van den Oord, A.; Li, Y.; and Vinyals, O. 2018.  Representation learning with contrastive predictive coding.  _arXiv_ , 1807.03748.
  * Wang et al. (2021) Wang, B.; Shen, T.; Long, G.; Zhou, T.; Wang, Y.; and Chang, Y. 2021.  Structure-augmented text representation learning for efficient knowledge graph completion.  In _WWW_ , 1737–1748.
  * Wang, Ma, and Gao (2020) Wang, Y.; Ma, F.; and Gao, J. 2020.  Efficient knowledge graph validation via cross-graph representation learning.  In _CIKM_ , 1595–1604.
  * Xie et al. (2018) Xie, R.; Liu, Z.; Lin, F.; and Lin, L. 2018.  Does William Shakespeare really write Hamlet? Knowledge representation learning with confidence.  In _AAAI_ , 4954–4961.
  * Yang et al. (2015) Yang, B.; tau Yih, S. W.; He, X.; Gao, J.; and Deng, L. 2015.  Embedding entities and relations for learning and inference in knowledge bases.  In _ICLR_.
  * Yang et al. (2022) Yang, C.; An, Z.; Cai, L.; and Xu, Y. 2022.  Mutual contrastive learning for visual representation learning.  In _AAAI_ , 3045–3053.
  * Yao, Mao, and Luo (2019) Yao, L.; Mao, C.; and Luo, Y. 2019.  KG-BERT: BERT for knowledge graph completion.  _arXiv_ , 1909.03193.
  * Zhang et al. (2022) Zhang, Q.; Dong, J.; Duan, K.; Huang, X.; Liu, Y.; and Xu, L. 2022.  Contrastive knowledge graph error detection.  In _CIKM_ , 2590–2599.
  * Zhao and Liu (2019) Zhao, Y.; and Liu, J. 2019.  SCEF: A support-confidence-aware embedding framework for knowledge graph refinement.  _arXiv_ , 1902.06377.