#  CrysFormer: Protein Structure Prediction via 3d Patterson Maps and Partial
Structure Attention

Chen Dun*
Rice University
<EMAIL_ADDRESS>
&Qiutai Pan*
Rice University
<EMAIL_ADDRESS>
&Shikai Jin
Rice University
<EMAIL_ADDRESS>
&Ria Stevens
Rice University
<EMAIL_ADDRESS>
&Mitchell D. Miller
Rice University
<EMAIL_ADDRESS>
&George N. Phillips, Jr.
Rice University
<EMAIL_ADDRESS>
&Anastasios Kyrillidis
Rice University
<EMAIL_ADDRESS>

###### Abstract

Determining the structure of a protein has been a decades-long open question.
A protein’s three-dimensional structure often poses nontrivial computation
costs, when classical simulation algorithms are utilized. Advances in the
transformer neural network architecture –such as AlphaFold2– achieve
significant improvements for this problem, by learning from a large dataset of
sequence information and corresponding protein structures. Yet, such methods
only focus on sequence information; other available prior knowledge, such as
protein crystallography and partial structure of amino acids, could be
potentially utilized. To the best of our knowledge, we propose the first
transformer-based model that directly utilizes protein crystallography and
partial structure information to predict the electron density maps of
proteins. Via two new datasets of peptide fragments (2-residue and 15-residue)
, we demonstrate our method, dubbed CrysFormer, can achieve accurate
predictions, based on a much smaller dataset size and with reduced computation
costs.

**footnotetext: Authors contributed equally.

##  1 Introduction

Proteins, the biological molecular machines, play a central role in the
majority of cellular processes (Tanford & Reynolds, 2004). The investigation
of a protein’s structure is a classic challenge in biology, given that its
function is dictated by its specific conformation. Proteins comprise long
chains of linked, relatively small organic molecules called amino acids, with
a set of twenty of them considered as standard. However, these underlying
polypeptide chains fold into complex three-dimensional structures, as well as
into larger assemblies thereof. Consequently, biologists aim to establish a
standardized approach for experimentally determining and visualizing the
overall structure of a protein at a low cost.

In the past decades, there have been three general approaches to the protein
structure problem: $i)$ ones that rely on physical experimental measurements,
such as X-ray crystallography, NMR, or cryo-electron microscopy; see (Drenth,
2007) for more details; $ii)$ protein folding simulation tools based on
thermodynamic or kinetic simulation of protein physics (Brini et al., 2020;
Sippl, 1990); and, $iii)$ evolutionary programs based on bioinformatics
analysis of the evolutionary history of proteins (Šali & Blundell, 1993; Roy
et al., 2010).

Recent advances in machine learning (ML) algorithms have inspired a fourth
direction which is to train a deep neural network model on a combination of a
large-scale protein structure data set (i.e., the Protein Data Bank (wwPDB
consortium, 2019)) and knowledge of the amino acid sequences of a vast number
of homologous proteins, to directly predict the protein structure from the
protein’s amino acid sequence. Recent research projects –such as Alphafold2
(Jumper et al., 2021)– further show that, with co-evolutionary bioinformatic
information (e.g., multiple sequence alignments), deep learning can achieve
highly accurate predictions in most cases.

Our hypothesis and contributions. While it is true that computational methods
of predicting structures without experimentally confirming data are improving,
they are not yet complete –in terms of the types of structures that can be
predicted– and suffer from lack of accuracy in many of the details
(Terwilliger et al., 2023). X-ray crystallographic data continues to be a gold
standard for critical details describing chemical interactions of proteins.

Having a robust and accurate way of going directly from an X-ray diffraction
pattern to a solved structure would be a strong contribution to the field of
X-ray crystallography. Such approaches are missing from the literature, with
the exception of Pan et al. (2023), a recent effort on the same problem based
on residual convolutional autoencoders.

Here, we present the first transformer-based model that utilizes protein
crystallography and partial structure information to directly predict the
electron density maps of proteins, going one step beyond such recent
approaches. While not yet ready to solve real problems, we demonstrate success
on a simplified problem. As a highlight, using a new dataset of small peptide
fragments of variable unit cell sizes –a byproduct of this work– we
demonstrate that our method, named CrysFormer, can achieve more accurate
predictions than state of the art (Pan et al., 2023) with less computations.

Some of our findings and contributions are:

  * •

CrysFormer is able to process the global information in Patterson maps to
infer electron density maps; to the best of our knowledge, along with Pan et
al. (2023), these are the first works to attempt this setting.

  * •

CrysFormer can incorporate “partial structure” information, when available; we
also show that such information could be incorporated in existing solutions
that neglected this feature, like the convolutional U-Net-based architectures
in Pan et al. (2023). However, the CrysFormer architecture still leads to
better reconstructions.

  * •

In practice, CrysFormer achieves a significant improvement in prediction
accuracy in terms of both Pearson coefficient and mean phase error, while
requiring both a smaller number of epochs to converge and less time taken per
epoch.

  * •

This work introduces a new dataset of variable-cell dipeptide fragments, where
all of the input Patterson and output electron density maps were derived from
the Protein Databank (PDB) (wwPDB consortium, 2019), solved by X-ray
Crystallography. We will make this dataset publicly available.

##  2 Problem Setup and Related Work

X-ray crystallography and the crystallographic phase problem. X-ray
crystallography has been the most commonly used method to determine a
protein’s electron density map111The electron density is a measure of the
probability of an electron being present just around a particular point in
space; a complete electron density map can be used to obtain a molecular model
of the unit cell. for over 100 years (Lattman & Loll, 2008). However, there is
an open question, called the crystallographic phase problem, that prevents
researchers from utilizing it to predict true structures/electron density
maps.

In review, each spot (known as a reflection) in an X-ray crystallography
diffraction pattern is denoted by three indices $h,k,l$, known as Miller
indices (Ashcroft & Mermin, 2022). These correspond to sets of parallel planes
within the protein crystal’s unit cell that contribute to producing the
reflections. The set of possible $h,k,l$ values is determined by the radial
extent of the observed diffraction pattern. Any reflection has an underlying
mathematical representation, known as a structure factor, dependent on the
locations and scattering factors of all the atoms within the crystal’s unit
cell. In math:

$\displaystyle F(h,k,l)=\sum_{j=1}^{n}f_{j}\cdot e^{2\pi
i(hx_{j}+ky_{j}+lz_{j})},$ (1)

where the scattering factor and location of atom $j$ are $f_{j}$ and
$(x_{j},y_{j},z_{j})$, respectively.

A structure factor $F(h,k,l)$ has both an amplitude and a phase component
(denoted by $\phi$) and thus can be considered a complex number. Furthermore,
suppose we knew both components of the structure factors corresponding to all
of the reflections within a crystal’s diffraction pattern. Then, in order to
produce an accurate estimate of the electron density at any point $(x,y,z)$
within the crystal’s unit cell, we would only need to take a Fourier transform
of all of these structures, as in:

$\displaystyle\rho(x,y,z)=\tfrac{1}{V}\cdot\sum_{h,k,l}|F(h,k,l)|\cdot
e^{-2\pi i(hx+ky+lz-\phi(h,k,l))},$ (2)

where $V$ is the volume of the unit cell. The amplitude $|F(h,k,l)|$ of any
structure factor is easy to determine, as it is simply proportional to the
square root of the measured intensity of the corresponding reflection.
However, it is impossible to directly determine the phase $\phi(h,k,l)$ of a
structure factor, and this is what is well-known as the crystallographic phase
problem (Lattman & Loll, 2008).

Solving the phase problem. Various methods have been developed to solve the
crystallography phase problem. The three commonly used methods are isomorphous
replacement, anomalous scattering, and molecular replacement (Lattman & Loll,
2008; Jin et al., 2020). Also, what is known as direct methods have been
successful for small molecules that diffract to atomic resolution, but they
rarely work for protein crystallography, due to the difficulty of resolving
atoms as separate objects. Alternative methods have been developed to solve
the phase problem based on intensity measurements alone, known as phase
retrieval (Guo et al., 2021; Kappeler et al., 2017; Rivenson et al., 2018).
However, these methods have not been widely used in X-ray crystallography,
because they assume different sampling conditions or were designed for non-
crystallographic fields of physics. The iterative non-convex Gerchberg–Saxton
algorithm (Fienup, 1982; Zalevsky et al., 1996) is a well-known example of
such methods, but requires more measurements than is available in
crystallography.

Although adaptations of the Gerchberg–Saxton algorithm have been proposed for
crystallography-like settings, they have not been used to solve the phase
problem except in special cases where crystals have very high solvent content
(He & Su, 2015; He et al., 2016; Kingston & Millane, 2022). More recently,
Candes et al. (2013) introduced the Phaselift method, a convex, complex
semidefinite programming approach, and Candes et al. (2015) the Wirtinger flow
algorithm (Candes et al., 2015), a non-convex phase retrieval method; both
these methods have not been applied practically, due to their computationally
intensive nature.

##  3 CrysFormer: Using 3d Maps and Partial Structure Attention

Inspired by (Hurwitz, 2020), we rely on deep learning solutions to directly
predict the electron density map of a protein. Later in the text, we
demonstrate that such a data-centric method achieves both better accuracy and
reduced computational cost.

The Patterson function. We utilize the Patterson function (Patterson, 1934), a
simplified variation of the Fourier transform from structure factors to
electron density, in which all structure factor amplitudes are squared, and
all phases are set to zero (i.e., ignored), as in:

$\displaystyle p(u,v,w)=\tfrac{1}{V}\cdot\sum_{h,k,l}|F(h,k,l)|^{2}\cdot
e^{-2\pi i(hu+kv+lw)}.$ (3)

It is important to note the Patterson map can be directly obtained from raw
diffraction data without the need for additional experiments, or any other
information.

Due to the discrete size of the input and output layers in deep learning
models, we can discretize and reformulate the electron density map –and its
corresponding Patterson map– as follows: Suppose the electron density map of a
molecule in interest is discretized into a $N_{1}\times N_{2}\times N_{3}$ 3d
grid. The electron density map can then be denoted as
$\mathbf{e}\in\mathbb{R}^{N_{1}\times N_{2}\times N_{3}}$. The Patterson map
is then formulated as follows, where $\odot$ means matrix element-wise
multiplication:

$\displaystyle\mathbf{p}=\Re\left(\mathcal{F}^{-1}\left(\mathcal{F}(\mathbf{e})\odot\mathcal{F}(\widehat{\mathbf{e}})\right)\right)\approx\Re\left(\mathcal{F}^{-1}\left(|\mathcal{F}(\mathbf{e})|^{2}\right)\right).$

Breaking down the above expression,
$\mathcal{F}(\mathbf{e})\odot\mathcal{F}(\widehat{\mathbf{e}})\approx|\mathcal{F}(\mathbf{e})|^{2}$
denotes only the magnitude part of the complex signals, as measured through
the Fourier transform of the input signal $\mathbf{e}$. Here,
$\widehat{\mathbf{e}}$ denotes an inverse-shifted version of $\mathbf{e}$,
where its entries follow the shifted rule as in
$\widehat{e}_{i,j,k}=e_{N-i,N-j,N-k}$.

Using deep learning. We follow a data-centric approach and train a deep
learning model, abstractly represented by $g(\bm{\theta},\cdot)$, such that
given a Patterson map $\mathbf{p}$ as input, it generates an estimate of an
electron density map, that resembles closely the true map $\mathbf{e}$.
Formally, given a data distribution $\mathcal{D}$ and
$\left\\{\mathbf{p}_{i},\mathbf{e}_{i}\right\\}_{i=1}^{n}\sim\mathcal{D}$,
where $\mathbf{p}_{i}\in\mathbb{R}^{N_{1}\times N_{2}\times N_{3}}$ is the
Patterson map that corresponds to the true data electron density map,
$\mathbf{e}_{i}\in\mathbb{R}^{N_{1}\times N_{2}\times N_{3}}$, deep learning
training aims in finding $\bm{\theta}^{\star}$ as in:

$\displaystyle\bm{\theta}^{\star}$
$\displaystyle=\underset{\bm{\theta}}{\operatorname*{arg\,min}}~{}\left\\{\mathcal{L}(\bm{\theta}):=\tfrac{1}{n}\sum_{i=1}^{n}\ell(\bm{\theta};~{}g,\\{\mathbf{p}_{i},\mathbf{e}_{i}\\})=\tfrac{1}{n}\sum_{i=1}^{n}\|g(\bm{\theta},\mathbf{p}_{i})-\mathbf{e}_{i}\|_{2}^{2}\right\\}.$

Since we have a regression problem, we use mean squared error as the loss
function $\mathcal{L}(\bm{\theta})$.

Using partial protein structures. Due to the well-studied structure of amino
acids, we aim to optionally utilize standardized partial structures to aid
prediction, when they are available. For example, let
$\mathbf{u}^{j}_{i}\in\mathbb{R}^{N_{1}\times N_{2}\times N_{3}}$ be the known
standalone electron density map of the $j$-th amino acid of the $i$-th protein
sample, in a standardized conformation. Abstractly, we then aim to optimize:

$\displaystyle\bm{\theta}^{\star}$
$\displaystyle=\underset{\bm{\theta}}{\operatorname*{arg\,min}}~{}\left\\{\mathcal{L}(\bm{\theta}):=\tfrac{1}{n}\sum_{i=1}^{n}\ell(\bm{\theta};~{}g,\\{\mathbf{p}_{i},\mathbf{e}_{i},\mathbf{u}^{j}_{i}\\})=\tfrac{1}{n}\sum_{i=1}^{n}\|g(\bm{\theta},\mathbf{p}_{i},\mathbf{u}^{j}_{i})-\mathbf{e}_{i}\|_{2}^{2}\right\\}.$

Challenges and Design Principles. We face the difficult learning problem to
infer electron density maps $\mathbf{e}$ from Patterson maps $\mathbf{p}$,
which involves Fourier transformations. These transformations can be
intuitively considered as transforming local information to global
information, which is rare in common deep model use cases. Secondly, it is
nontrivial to incorporate the partial structure density maps
$\mathbf{u}^{j}_{i}$ to aid prediction. Thirdly, the 3d data format of both
our inputs and outputs often increases substantially the computational
requirements. Finally, since part of our contributions is novel datasets on
this problem, we need to be data efficient due to the expensive dataset
creation cost. Thus, the main design principles for our model can be
summarized as:

  * •

Design Principle #1: Be able to process the global information in Patterson
maps to correctly infer the corresponding electron density maps;

  * •

Design Principle #2: Be able to incorporate partial structure information,
when available;

  * •

Design Principle #3: Learn to fulfill the above, with reduced computational
and data-creation costs.

Gap in current knowledge. As an initial attempt, the well-established
convolution-based U-Net model (Ronneberger et al., 2015) could be utilized for
this task. This is the path followed in (Pan et al., 2023). However, classical
U-Nets cannot fulfill the design principles above, since: $i)$ they mostly
rely on local information within CNN layers; such a setup is not suitable when
Patterson maps are available, since the latter do not have meaningful local
structures. $ii)$ It is not clear (or, at best, non-trivial) to incorporate
any partial protein structures prior information, since the latter is in a
different representation domain, compared to Patterson maps. Finally, $iii)$ a
large 3d U-Net model is computationally expensive and inefficient, due to the
3d filter convolution computation.

Our proposal: CrysFormer. We propose CrysFormer, a novel, 3d Transformer model
(Vaswani et al., 2017; Chen et al., 2021) with a new self-attention mechanism
to process Patterson maps and partial protein structures, to directly infer
electron density maps with reduced costs.

Inspired by recent research on the potential connection between Fourier
transforms and the self-attention mechanism, found in the Transformer model
(Lee-Thorp et al., 2022), CrysFormer captures the global information in
Patterson maps and “translates” it into correct electron density map
predictions, via our proposed self-attention mechanism (Design Principle #1).
CrysFormer does not need an encoder-decoder structure (Vaswani et al., 2017)
and artificial information bottlenecks (Cheng et al., 2019) –as in the U-Net
architecture– to force the learning of global information.

By definition, CrysFormer is able to handle additional partial structure
information, which comes from a different domain than the Patterson maps
(Design Principle #2; more details below).

Finally, by using efficient self-attention between 3d image patches, we can
significantly reduce the overall computation cost. Detaching our model from an
encoder-decoder architecture further reduces the required depth of the model
and, thus, the overall training cost (Design Principle #3).

The architecture of the CrysFormer. We follow ideas of a 3d visual Transformer
(Chen et al., 2021) by partitioning the whole input 3d Patterson map
$\mathbf{p}_{i}\in\mathbb{R}^{N_{1}\times N_{2}\times N_{3}}$ input into a set
of smaller 3d patches. We embed them into one-dimensional “word tokens”, and
feed them into a multi-layer, encoder-only Transformer module. If partial
structures $\mathbf{u}^{j}_{i}$ are also available, we will partition them
into 3d patches and embed them into additional tokens that are sent to each
self-attention layer. This way, the tokens in each layer can also “attend” the
election density of partial structures, as a reference for final global
electron density map predictions. Finally, we utilize a 3d convolutional layer
to transform “word-tokens” back into a 3d electron density map.222We also
utilize 3d convolutional layer(s) at the very beginning of the execution to
expand the number of channels of the Patterson map (and potentially partial
structure) inputs. See Figure 1.

Figure 1: Abstract depiction of the Crysformer, which utilizes a one-way
attention mechanism (red and purple arrows) to incorporate the partial
structure information. The tokens from the additional partial structure all
come from initial 3d CNN embedding and are not passed to the next layer.

Mathematically, we report the following: The first part is the preprocessing
and partitioning of input Patterson maps $\mathbf{p}$ and additional partial
structures $\mathbf{u}^{j}$ into 3d patches of size $d_{1}\times d_{2}\times
d_{3}$. We embed those patches into one-dimensional tokens with dimension
$d_{t}$, using of a small MLP, and add them with a learned positional
embedding; this holds for both Patterson maps and structures, as below:

Patterson maps $\mathbf{p}$

$\displaystyle\mathbf{X}^{0}$
$\displaystyle=\texttt{3DCNN}_{\mathbf{W}_{c}}(\mathbf{p})\in\mathbb{R}^{c\times
N_{1}\times N_{2}\times N_{3}}$ $\displaystyle\mathbf{X}^{0}$
$\displaystyle=\texttt{Partition}(\mathbf{X}^{0})\in\mathbb{R}^{\frac{N_{1}}{d_{1}}\times\frac{N_{2}}{d_{2}}\times\frac{N_{3}}{d_{3}}\times(cd_{1}d_{2}d_{3})}$
$\displaystyle\mathbf{X}^{0}$
$\displaystyle=\texttt{Flatten}(\mathbf{X}^{0})\in\mathbb{R}^{\frac{N_{1}N_{2}N_{3}}{d_{1}d_{2}d_{3}}\times(cd_{1}d_{2}d_{3})}$
$\displaystyle\mathbf{X}^{0}$
$\displaystyle=\texttt{MLP}_{\mathbf{W}_{c}}(\mathbf{X}^{0})\in\mathbb{R}^{\frac{N_{1}N_{2}N_{3}}{d_{1}d_{2}d_{3}}\times
d_{t}}$ $\displaystyle\mathbf{X}^{0}$
$\displaystyle=\mathbf{X}^{0}+\texttt{PosEmbedding}(\tfrac{N_{1}N_{2}N_{3}}{d_{1}d_{2}d_{3}})$

Partial structures $\mathbf{u}^{j}$

$\\!\begin{aligned}
\mathbf{U}^{j}&=\texttt{3DCNN}_{\mathbf{W}_{p}}(\mathbf{u}^{j})\in\mathbb{R}^{c\times
N_{1}\times N_{2}\times N_{3}}\\\
\mathbf{U}^{j}&=\texttt{Partition}(\mathbf{U}^{j})\in\mathbb{R}^{\frac{N_{1}}{d_{1}}\times\frac{N_{2}}{d_{2}}\times\frac{N_{3}}{d_{3}}\times(cd_{1}d_{2}d_{3})}\\\
\mathbf{U}^{j}&=\texttt{Flatten}(\mathbf{U}^{j})\in\mathbb{R}^{\frac{N_{1}N_{2}N_{3}}{d_{1}d_{2}d_{3}}\times(cd_{1}d_{2}d_{3})}\\\
\mathbf{U}^{j}&=\texttt{MLP}_{\mathbf{W}_{p}}(\mathbf{U}^{j})\in\mathbb{R}^{\tfrac{N_{1}N_{2}N_{3}}{d_{1}d_{2}d_{3}}\times
d_{t}}\\\
\mathbf{U}^{j}&=\mathbf{U}^{j}+\texttt{PosEmbedding}(\tfrac{N_{1}N_{2}N_{3}}{d_{1}d_{2}d_{3}})\end{aligned}$
---

As shown in Figure 1, we design an efficient attention mechanism such that
$i)$ only tokens from Patterson maps attend tokens from the partial
structures; $ii)$ the tokens from the additional partial structures are not
passed to the next layer. This is based on that the partial structure electron
density information should be used by the model as a stable reference to
attend to in each layer.

This one-way attention also greatly reduces the overall communication cost. In
particular, let the token sequence length be
$S=\tfrac{N_{1}N_{2}N_{3}}{d_{1}d_{2}d_{3}}$ and let $d_{h}$ denote the
dimension of the attention head. Assuming we have $H$ attention heads and $L$
layers, CrysFormer uses the following attention mechanism:

$\displaystyle\mathbf{U}$
$\displaystyle=\texttt{Concat}_{j=1}^{J}(\mathbf{U}^{j})\in\mathbb{R}^{(SJ)\times
d_{t}}$ $\displaystyle\mathbf{A}^{h}$
$\displaystyle=\texttt{Softmax}\left((\mathbf{W}^{h}_{q}\mathbf{X}^{\ell})^{\top}(\texttt{Concat}(\mathbf{W}^{h}_{k}\mathbf{X}^{\ell},\mathbf{W}^{h}_{k^{\prime}}\mathbf{U})\right)\in\mathbb{R}^{S\times(S+SJ)};$
$\displaystyle\mathbf{\widehat{V}}^{h}$
$\displaystyle=\mathbf{A}^{h}\left(\texttt{Concat}(\mathbf{W}^{h}_{v}\mathbf{X}^{\ell},\mathbf{W}^{h}_{v^{\prime}}\mathbf{U})\right)\in\mathbb{R}^{S\times
d_{h}};$ $\displaystyle\mathbf{O}$
$\displaystyle=\mathbf{W}_{o}\texttt{Concat}\left(\mathbf{\widehat{V}}^{0},~{}\mathbf{\widehat{V}}^{1},\dots,~{}\mathbf{\widehat{V}}^{H}\right)\in\mathbb{R}^{S\times
d_{t}};$ $\displaystyle\mathbf{X}^{\ell+1}$
$\displaystyle=\mathbf{W}_{\text{ff2}}(\texttt{ReLU}(\mathbf{W}_{\text{ff1}}\mathbf{O})),$

where, omitting the layer index, $\mathbf{W}^{h}_{q}$, $\mathbf{W}^{h}_{k}$,
$\mathbf{W}^{h}_{v}$ are the trainable query, key, and value projection
matrices of the $h$-th attention head for tokens from the Patterson map, and
$\mathbf{W}^{h}_{k^{\prime}}$, $\mathbf{W}^{h}_{v^{\prime}}$ are the
corresponding matrices for tokens from the partial structure, each with
dimension $d_{h}$. Further, $\mathbf{W}_{\text{ff1}}$ and
$\mathbf{W}_{\text{ff2}}$ are the trainable parameters of the fully-connected
layers. We omit skip connections and layer normalization modules just to
simplify notation, but these are included in practice.

As a final step, we transform the output embedding back to a 3d electron
density map, as follows:

$\displaystyle g(\bm{\theta},\mathbf{p})$
$\displaystyle=\texttt{tanh}(\texttt{3DCNN}_{\mathbf{W}_{o}}(\texttt{Rearrange}(\texttt{MLP}(\mathbf{X}^{L}))))\in\mathbb{R}^{N_{1}\times
N_{2}\times N_{3}},$

and, as stated previously, we use as our loss function the standard mean
squared error loss.

##  4 New Datasets

We generate datasets of protein fragments, where input Patterson and output
electron density maps are derived from Protein Databank (PDB) entries of
proteins solved by X-ray Crystallography (wwPDB consortium, 2019). We start
from a curated basis of $\sim 24,000$ such protein structures. Then from a
random subset of about half of these structures, we randomly select and store
segments of adjacent amino acid residues. These examples are consisted of
dipeptides (two residues) and 15-residues, leading to two datasets that we
introduce with this work. The latter dataset contains 15 residues, where at
most 3 residues could be shared between different examples. Using the pdbfixer
Python API (Eastman et al., 2017), we remove all examples that either contain
nonstandard residues or have missing atoms from our initial set. We also apply
a few standardized modifications.

For our dipeptide dataset, we then iteratively expand the unit cell dimensions
for each example, starting from the raw $\max-\min$ ranges in each of the
three axis directions, attempting to create a minimal-size unit cell where the
minimum atomic contact is at least $2.75$ Angstroms (Å).333An Angstrom is a
metric unit of length equal to $10^{-10}$m. For our 15-residue dataset, we
instead place atoms in fixed unit cells of size $41$ Å x $30$ Å x $24$ Å to
simplify the now much harder problem. After this, all examples that still
contain atomic contacts of less than $2.75$ Å are discarded. The examples are
then reoriented via a reindexing operation, such that the first axis is always
the longest and the third axis is always the shortest.

One issue leading to potential ambiguity in interpreting Patterson maps is
their invariance to translation of the entire corresponding electron density
(Hurwitz, 2020). To tackle this, we center all atomic coordinates such that
the center of mass is in the center of the corresponding unit cell. This means
that our model’s predicted electron densities would always be more or less
centered in the unit cell. We note that this is also the case for the majority
of actual protein crystals.

Structure factors for each remaining example, as well as those for the
corresponding partial structures for each of the present amino acids, are
generated using the gemmi sfcalc program (Wojdyr, 2022) to a resolution of
$1.5$ Å. An electron density and Patterson map for each example are then
obtained from those structure factors with the fft program of the CCP4 program
suite (Read & Schierbeek, 1988; Winn et al., 2011); partial structure
densities are obtained in the same manner. We specify a grid oversampling
factor of $3.0$, resulting in a $0.5$ Å grid spacing in the produced maps. All
these maps are then converted into PyTorch tensors. We then normalize the
values in each of the tensors to be in the range $[-1,~{}1]$. Since, in our
PyTorch implementation, all examples within a training batch are of the same
size, we remove all examples from the tensor-size bins containing fewer
examples than a specified minimum batch size.

##  5 Experiments

Baselines. There are no readily available off-the-self solutions for our
setting, as our work is one of the first of this kind. As our baseline, we use
a CNN-based U-Net model (Pan et al., 2023); this architecture is widely used
in image transformation tasks (Ronneberger et al., 2015; Yan et al., 2021).

For comparison, we have further enhanced this vanilla U-Net with $i)$
additional input channels to incorporate the partial structure information,
despite being evidently unsound; and $ii)$ a refining model procedure, which
retrains the U-Net using previous model predictions as additional input
channels. Both of these extensions are shown to greatly improve the
performance of the vanilla U-Net. We refer the reader to the appendix for more
details on our baseline model architecture.

Metrics. During testing, we calculate the Pearson correlation coefficient
between the ground truth targets $\mathbf{e}$ and model predictions
$g(\bm{\theta},\mathbf{p})$; the larger this coefficient is, the better. Let
us denote a model prediction as $\mathbf{e^{\prime}}$. We define
$\bar{\mathbf{e}}=\frac{1}{N_{1}N_{2}N_{3}}\sum_{i,j,k}\mathbf{e}_{i,j,k}$ and
$\bar{\mathbf{e}}^{\prime}=\frac{1}{N_{1}N_{2}N_{3}}\sum_{i,j,k}\mathbf{e}^{\prime}_{i,j,k}$.
Then, the Pearson correlation coefficient between $\mathbf{e}$ and
$\mathbf{e^{\prime}}$ is as below:

$\texttt{PC}(\mathbf{e},\mathbf{e}^{\prime})=\frac{\sum_{i,j,k=1}^{N_{1},N_{2},N_{3}}(\mathbf{e}^{\prime}_{i,j,k}-\bar{\mathbf{e}}^{\prime})(\mathbf{e}_{i,j,k}-\bar{\mathbf{e}})}{\rule{0.0pt}{16.00008pt}\sqrt{\sum_{i,j,k=1}^{N_{1},N_{2},N_{3}}(\mathbf{e}^{\prime}_{i,j,k}-\bar{\mathbf{e}}^{\prime})+\epsilon}\cdot\>\sqrt{\sum_{i,j,k=1}^{N_{1},N_{2},N_{3}}(\mathbf{e}_{i,j,k}-\bar{\mathbf{e}})+\epsilon}},$
(4)

where $\epsilon$ is a small constant to prevent division by zero. To
demonstrate how well our methods solve the phase problem, we also perform
phase error analysis on our models’ final post-training predictions using the
cphasematch program of the CCP4 program suite (Cowtan, 2011). We report the
mean phase errors of our predictions in degrees, as reported by cphasematch,
where a smaller phase error is desirable. Finally, we compare the convergence
speed and computation cost of both methods.

Method | Mean $\texttt{PC}(\mathbf{e},\mathbf{e}^{\prime})$ | Mean Phase Error | Epochs | Time per epoch (mins.)
---|---|---|---|---
U-Net (Pan et al., 2023) | 0.735 | 67.40∘ | 50 | 28.93
U-Net+R (This work) | 0.775 | 58.67∘ | 90 | 29.06
U-Net+PS+R (This work) | 0.839 | 51.34∘ | 90 | 29.31
CrysFormer (This work) | 0.939 | 35.16∘ | 35 | 12.37
Table 1: CrysFormer versus baselines on the dipeptide dataset. U-Net+R refers
to adding the refining procedure to U-Net training; U-Net+PS+R refers to
adding further partial structures as additional channels.

Results on two-residues. A summary of our results on our dipeptide dataset,
which consisted of $1,894,984$ training and $210,487$ test cases, is provided
in Table 1. Overall, CrysFormer achieves a significant improvement in
prediction accuracy in terms of both the Pearson coefficient and phase error,
while requiring a shorter time (in epochs) to converge. CrysFormer also incurs
much less computation cost which results in significantly reduced wall clock
time per epoch.

We further visualize some of the predictions in Figure LABEL:fig:main,
comparing side by side those made by the baselines and the CrysFormer.
CrysFormer produces more accurate predictions in terms of both global and
local structures. This verifies our hypothesis that $i)$ the self-attention
mechanism can better capture the global information in Patterson maps, and
$ii)$ the removal of the U-Net’s encoder-decoder structure prevents loss of
information and improves the reproduction of finer details.

E.g., the top row of Figure LABEL:fig:main represents a class of examples
containing a large aromatic residue, Tryptophan. U-Net+R models consistently
produce poor predictions in this case, while the CrysFormer better handles
such residues. U-Net+PS+R shows that both providing additional input channels
and using the refining procedure improves results even for U-Net
architectures; yet, CrysFormer still provides better reconstruction. More
visualizations can be found in the appendix.

We further plot the calculated average mean phase errors of the predictions of
our models against reflection resolution, see left panel of Figure
LABEL:phase_error. The predictions made by CrysFormer have lower mean phase
error, compared to baselines. This means that the CrysFormer predictions, on
average, can reproduce better the general shape, as well as finer details of
the ground truth electron densities.

Finally, we generate a chart of the fraction of our models’ predictions for
which the calculated mean phase error is $<60^{\circ}$ at various ranges of
resolution. We consider such predictions to accurately reproduce the level of
detail specified by that resolution range. This is shown on the right panel in
Figure LABEL:phase_error. At all resolution ranges, CrysFormer predictions are
clearly better than that of the U-Net-based models. In particular, for
CrysFormer, we still have a majority of predictions with phase error
$<60^{\circ}$ even at the highest ranges of resolution.

Results on 15-residues. On our dataset of 15-residue examples, which consisted
of only $165,858$ training and $16,230$ test cases (less than one-tenth the
size of our dipeptide dataset), we trained for 80 epochs to a final average
test set Pearson correlation of about $0.747$. We then performed a refining
training run of 20 epochs, incorporating the original training run’s
predictions as additional input channels when training the CrysFormer, and
obtained an improved average test set Pearson correlation of about $0.77$ and
phase error of about $67.66$. On both of these runs, we used the Nyström
approximate attention mechanism (Xiong et al., 2021) when incorporating our
partial structure information to reduce time and space costs. Even still, each
training epoch still took about $6.28$ hours to complete. Thus due to time
considerations, we decided not to attempt to train a U-Net on this dataset for
purposes of comparison.

We provide visualizations of some model predictions in Figure LABEL:pred15;
more can again be found in the appendix. We also plot the average mean phase
errors of the predictions of our models against reflection resolution, as well
as the fraction of our models’ predictions for which the calculated mean phase
error is $<60^{\circ}$ at various ranges of resolution in Figure
LABEL:phase_error15. These results show that this is a more difficult dataset
with reduced sample size; yet CrysFormer predictions tend to accurately
reproduce details of the desired electron densities.

We used the Autobuild program within the PHENIX suite (Terwilliger et al.,
2008; Liebschner et al., 2019) to perform automated model building and
crystallographic refinement on a randomly selected subset of $302$ test set
predictions after the refining training run. We found that $281$ out of $302$
($\sim 93\%$) refined to a final atomic model with a crystallographic
$R$-factor of less than $0.38$, indicating success, when solvent flattening
was applied. Without solvent flattening, $258$ out of $302$ ($\sim 85\%$)
refined to such an $R$-factor (performing solvent flattening is known to be
especially effective for unit cells with high solvent content, i.e. a large
amount of empty space around the atoms). Figure LABEL:autobuild15 shows these
results as scatterplots; clearly only a small fraction of the subset of
predictions did not refine successfully. And even if no refinement was
performed at all, and instead an atomic model was repeatedly fit to our
predicted electron densities, we found that $229$ out of $302$ ($\sim 76\%$)
of the best such atomic models still had a crystallographic $R$-factor of less
than $0.38$.

Furthermore, after automatic map interpretation using the autobuilding
routines in shelxe (Usón & Sheldrick, 2018) to obtain a poly-alanine chain
from each of the $16230$ test set predictions, we found that almost $74\%$ of
the resulting models had calculated amplitudes with a Pearson correlation of
at least $0.25$ to the true underlying data. Historical results indicate that
further refinement would very likely produce a "correct" model if the initial
poly-alanine model has at least such a correlation.

##  6 Discussion

We have shown that CrysFormer outperforms state of the art models for
predicting electron density maps from corresponding Patterson maps in all
metrics on a newly introduced dataset (dipeptide). Overall, CrysFormer
requires fewer epochs to reasonably converge and has a smaller computational
footprint. Furthermore, our “refining” procedure greatly improves training for
the vanilla U-Net architecture on our dipeptide dataset, as well as for
training CrysFormer on our both dipeptide and 15-residues dataset.

Limitations and next steps. Following successful results on our initial
15-residue dataset, we also suggest training our model on variable unit cells
at that problem size as future work. Eventually, we also prefer handling
variable cell angles as well, moving beyond the orthorhombic crystal system.
We will explore changing the formulation of our partial structures to have
more than one amino acid residue in a structure, as having each partial
structure representing only a single residue may no longer be reasonable, both
computationally and from a practical perspective.

Broader Impacts. Solving the crystallographic phase problem for proteins would
dramatically reduce the time and expense of determining a new protein
structure, especially if there are no close homologs already in the Protein
Data Bank. There exist some methods that sometimes work under special
conditions Jiang et al. (2018), or that work sometimes but only at very low
resolutions David & Subbiah (1994). The recent line of work on AlphaFold
Jumper et al. (2021); Tunyasuvunakool et al. (2021) definitely helps in these
problems; we note though that this is true mostly in cases where reliable
predictions are possible due to strong homologs and/or extensive sequence
data.

## References

  * Ashcroft & Mermin (2022) Neil W Ashcroft and N David Mermin.  _Solid state physics_.  Cengage Learning, 2022.
  * Brini et al. (2020) Emiliano Brini, Carlos Simmerling, and Ken Dill.  Protein storytelling through physics.  _Science_ , 370(6520):eaaz3041, 2020.
  * Candes et al. (2013) Emmanuel J Candes, Thomas Strohmer, and Vladislav Voroninski.  Phaselift: Exact and stable signal recovery from magnitude measurements via convex programming.  _Communications on Pure and Applied Mathematics_ , 66(8):1241–1274, 2013.
  * Candes et al. (2015) Emmanuel J Candes, Xiaodong Li, and Mahdi Soltanolkotabi.  Phase retrieval via wirtinger flow: Theory and algorithms.  _IEEE Transactions on Information Theory_ , 61(4):1985–2007, 2015.
  * Chen et al. (2021) Junyu Chen, Yufan He, Eric C. Frey, Ye Li, and Yong Du.  Vit-v-net: Vision transformer for unsupervised volumetric medical image registration.  _arXiv_ , .2104.06468, 2021.  URL https://doi.org/10.48550/arXiv.2104.06468.
  * Cheng et al. (2019) Hao Cheng, Dongze Lian, Shenghua Gao, and Yanlin Geng.  Utilizing information bottleneck to evaluate the capability of deep neural networks for image classification.  _Entropy_ , 21(5), 2019.  ISSN 1099-4300.  doi: 10.3390/e21050456.  URL https://www.mdpi.com/1099-4300/21/5/456.
  * Cowtan (2011) Kevin Cowtan.  cphasematch, 2011.  URL https://www.ccp4.ac.uk/html/cphasematch.html.
  * David & Subbiah (1994) P. R. David and S. Subbiah.  Low-resolution real-space envelopes: the application of the condensing-protocol approach to the ab initio macromolecular phase problem of a variety of examples.  _Acta Crystallographica Section D_ , 50(2):132–138, Mar 1994.  doi: 10.1107/S090744499301131X.  URL https://doi.org/10.1107/S090744499301131X.
  * Drenth (2007) Jan Drenth.  _Principles of protein X-ray crystallography_.  Springer Science & Business Media, 2007.
  * Eastman et al. (2017) Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle A Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al.  Openmm 7: Rapid development of high performance algorithms for molecular dynamics.  _PLoS computational biology_ , 13(7):e1005659, 2017.
  * Fienup (1982) J. R. Fienup.  Phase retrieval algorithms: a comparison.  _Appl. Opt._ , 21(15):2758–2769, Aug 1982.  doi: 10.1364/AO.21.002758.
  * Guo et al. (2021) Youming Guo, Yu Wu, Ying Li, Xuejun Rao, and Changhui Rao.  Deep phase retrieval for astronomical Shack–Hartmann wavefront sensors.  _Monthly Notices of the Royal Astronomical Society_ , 510(3):4347–4354, 12 2021.  ISSN 0035-8711.  doi: 10.1093/mnras/stab3690.
  * He & Su (2015) Hongxing He and Wu-Pei Su.  Direct phasing of protein crystals with high solvent content.  _Acta Crystallographica Section A_ , 71(1):92–98, Jan 2015.  doi: 10.1107/S2053273314024097.
  * He et al. (2016) Hongxing He, Hengrui Fang, Mitchell D Miller, George N Phillips Jr, and W-P Su.  Improving the efficiency of molecular replacement by utilizing a new iterative transform phasing algorithm.  _Acta Crystallographica Section A: Foundations and Advances_ , 72(5):539–547, 2016.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.  In _2015 IEEE International Conference on Computer Vision (ICCV)_ , pp. 1026–1034, New York, NY, USA, 2015. IEEE Press.  doi: 10.1109/ICCV.2015.123.
  * Helmenstine (2021) Anne Marie Helmenstine.  Amino acid chirality, 2021.  URL https://www.thoughtco.com/amino-acid-chirality-4009939.
  * Hu et al. (2018) Jie Hu, Li Shen, and Gang Sun.  Squeeze-and-excitation networks.  In _2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 7132–7141, New York, NY, USA, 2018. IEEE Press.  doi: 10.1109/CVPR.2018.00745.
  * Hurwitz (2020) David Hurwitz.  From patterson maps to atomic coordinates: Training a deep neural network to solve the phase problem for a simplified case.  _arXiv_ , 03 2020.
  * Jiang et al. (2018) Mengchao Jiang, Hongxing He, Yunpeng Cheng, and Wu-Pei Su.  Resolution dependence of an ab initio phasing method in protein x-ray crystallography.  _Crystals_ , 8(4), 2018.  ISSN 2073-4352.  doi: 10.3390/cryst8040156.  URL https://www.mdpi.com/2073-4352/8/4/156.
  * Jin et al. (2020) Shikai Jin, Mitchell D Miller, Mingchen Chen, Nicholas P Schafer, Xingcheng Lin, Xun Chen, GN Phillips, and PG Wolynes.  Molecular-replacement phasing using predicted protein structures from awsem-suite.  _IUCrJ_ , 7(6):1168–1178, 2020.
  * Jumper et al. (2021) John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie, Bernardino Romera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis.  Highly accurate protein structure prediction with alphafold.  _Nature_ , 596(7873):583–589, 2021.
  * Kappeler et al. (2017) Armin Kappeler, Sushobhan Ghosh, Jason Holloway, Oliver Cossairt, and Aggelos Katsaggelos.  Ptychnet: Cnn based fourier ptychography.  In _2017 IEEE International Conference on Image Processing (ICIP)_ , pp. 1712–1716, New York, NY, USA, 2017. IEEE Press.
  * Kingston & Millane (2022) Richard Lawrence Kingston and Rick P Millane.  A general method for directly phasing diffraction data from high-solvent-content protein crystals.  _IUCrJ_ , 9(5), 2022.
  * Lattman & Loll (2008) Eaton Lattman and Patrick Loll.  _Protein Crystallography_.  Johns Hopkins University Press, 2008.
  * Lee-Thorp et al. (2022) James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon.  Fnet: Mixing tokens with fourier transforms, 2022.
  * Liebschner et al. (2019) Dorothee Liebschner, Pavel V. Afonine, Matthew L. Baker, Gábor Bunkóczi, Vincent B. Chen, Tristan I. Croll, Bradley Hintze, Li-Wei Hung, Swati Jain, Airlie J. McCoy, Nigel W. Moriarty, Robert D. Oeffner, Billy K. Poon, Michael G. Prisant, Randy J. Read, Jane S. Richardson, David C. Richardson, Massimo D. Sammito, Oleg V. Sobolev, Duncan H. Stockwell, Thomas C. Terwilliger, Alexandre G. Urzhumtsev, Lizbeth L. Videau, Christopher J. Williams, and Paul D. Adams.  Macromolecular structure determination using x-rays, neutrons and electrons: recent developments in phenix.  _Acta Crystallogr._ , D75(10):861–877, Oct 2019.  doi: 10.1107/S2059798319011471.  URL https://doi.org/10.1107/S2059798319011471.
  * Pan et al. (2023) Tom Pan, Shikai Jin, Mitchell D Miller, Anastasios Kyrillidis, and George N Phillips.  A deep learning solution for crystallographic structure determination.  _IUCrJ_ , 10(4):487–496, 2023.
  * Patterson (1934) A. L. Patterson.  A fourier series method for the determination of the components of interatomic distances in crystals.  _Phys. Rev._ , 46:372–376, Sep 1934.  doi: 10.1103/PhysRev.46.372.
  * Read & Schierbeek (1988) RJ Read and AJ Schierbeek.  A phased translation function.  _Journal of Applied Crystallography_ , 21(5):490–495, 1988.
  * Rivenson et al. (2018) Yair Rivenson, Yibo Zhang, Harun Günaydın, Da Teng, and Aydogan Ozcan.  Phase recovery and holographic image reconstruction using deep learning in neural networks.  _Light: Science & Applications_, 7(2):17141–17141, 2018.
  * Ronneberger et al. (2015) Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18_ , pp. 234–241, 2015.
  * Roy et al. (2010) Ambrish Roy, Alper Kucukural, and Yang Zhang.  I-TASSER: a unified platform for automated protein structure and function prediction.  _Nature protocols_ , 5(4):725–738, 2010.
  * Šali & Blundell (1993) Andrej Šali and Tom L Blundell.  Comparative protein modelling by satisfaction of spatial restraints.  _Journal of molecular biology_ , 234(3):779–815, 1993.
  * Sippl (1990) Manfred J Sippl.  Calculation of conformational ensembles from potentials of mena force: an approach to the knowledge-based prediction of local structures in globular proteins.  _Journal of molecular biology_ , 213(4):859–883, 1990.
  * Tanford & Reynolds (2004) Charles Tanford and Jacqueline Reynolds.  _Nature’s Robots: A History of Proteins_.  Oxford University Press, 2004.
  * Terwilliger et al. (2008) Thomas C. Terwilliger, Ralf W. Grosse-Kunstleve, Pavel V. Afonine, Nigel W. Moriarty, Peter H. Zwart, Li-Wei Hung, Randy J. Read, and Paul D. Adams.  Iterative model building, structure refinement and density modification with the PHENIX AutoBuild wizard.  _Acta Crystallogr._ , D64(1):61–69, Jan 2008.  doi: 10.1107/S090744490705024X.  URL https://doi.org/10.1107/S090744490705024X.
  * Terwilliger et al. (2023) Thomas C. Terwilliger, Dorothee Liebschner, Tristan I. Croll, Christopher J. Williams, Airlie J. McCoy, Billy K. Poon, Pavel V. Afonine, Robert D. Oeffner, Jane S. Richardson, Randy J. Read, and Paul D. Adams.  Alphafold predictions are valuable hypotheses, and accelerate but do not replace experimental structure determination.  _bioRxiv_ , 2023.  doi: 10.1101/2022.11.21.517405.  URL https://www.biorxiv.org/content/early/2023/05/19/2022.11.21.517405.
  * Tunyasuvunakool et al. (2021) Kathryn Tunyasuvunakool, Jonas Adler, Zachary Wu, Tim Green, Michal Zielinski, Augustin Žídek, Alex Bridgland, Andrew Cowie, Clemens Meyer, Agata Laydon, et al.  Highly accurate protein structure prediction for the human proteome.  _Nature_ , 596(7873):590–596, 2021.
  * Usón & Sheldrick (2018) Isabel Usón and George M. Sheldrick.  An introduction to experimental phasing of macromolecules illustrated by SHELX; new autotracing features.  _Acta Crystallogr._ , D74(2):106–116, Feb 2018.  doi: 10.1107/S2059798317015121.  URL https://doi.org/10.1107/S2059798317015121.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Winn et al. (2011) Martyn D. Winn, Charles C. Ballard, Kevin D. Cowtan, Eleanor J. Dodson, Paul Emsley, Phil R. Evans, Ronan M. Keegan, Eugene B. Krissinel, Andrew G. W. Leslie, Airlie McCoy, Stuart J. McNicholas, Garib N. Murshudov, Navraj S. Pannu, Elizabeth A. Potterton, Harold R. Powell, Randy J. Read, Alexei Vagin, and Keith S. Wilson.  Overview of the CCP4 suite and current developments.  _Acta Crystallographica Section D_ , 67(4):235–242, Apr 2011.  doi: 10.1107/S0907444910045749.
  * Wojdyr (2022) Marcin Wojdyr.  Gemmi: A library for structural biology.  _Journal of Open Source Software_ , 7(73):4200, 2022.
  * wwPDB consortium (2019) wwPDB consortium.  Protein Data Bank: the single global archive for 3D macromolecular structure data.  _Nucleic Acids Research_ , 47(D1):D520–D528, 2019.  doi: 10.1093/nar/gky949.  URL https://doi.org/10.1093/nar/gky949.
  * Xiong et al. (2021) Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh.  Nyströmformer: A nyström-based algorithm for approximating self-attention.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 35(16):14138–14148, May 2021.  doi: 10.1609/aaai.v35i16.17664.  URL https://ojs.aaai.org/index.php/AAAI/article/view/17664.
  * Yan et al. (2021) Weidan Yan, Can Chen, and Dengyin Zhang.  U-net-based medical image segmentation algorithm.  In _13th International Conference on Wireless Communications and Signal Processing (WCSP)_ , pp. 1–5, 2021.  doi: 10.1109/WCSP52459.2021.9613447.
  * Zalevsky et al. (1996) Zeev Zalevsky, David Mendlovic, and Rainer G Dorsch.  Gerchberg–saxton algorithm applied in the fractional fourier or the fresnel domain.  _Optics Letters_ , 21(12):842–844, 1996.

##  Appendix A Appendix

##  Appendix B Model Architecture of Baseline U-Net model

Our U-net architecture can be divided into three phases. The Encoding Phase
consists of two 7x7x7 convolutional layers. The first has 25 output channels
while the second has 30 output channels, and both are followed by the standard
batch normalization and a ReLU activation. Then, a max pooling operation with
kernel size 2x2x2 and stride 2 is used to reduce the height, width, and depth
dimensions by a factor of 2.

The Learning Features Phase consists of a sequence of 7 residual blocks. Each
of these blocks consists of a 7x7x7 convolutional layer with 30 output
channels followed by batch normalization and ReLU activation, and then another
30-channel 7x7x7 convolutional layer with batch normalization but no
activation. A squeeze and excitation block Hu et al. (2018) occurs at this
point, applied with the channel dimension reduced by a factor of 2. Afterward,
the residual skip connection is applied, followed by another ReLU activation.
At the end of this phase, a naive upsampling operation is used to increase the
height, width, and depth dimensions by a factor of 2, thus restoring the
original dimensions.

The Decoding Phase consists of two 5x5x5 convolutional layers. The first has
25 output channels and is followed by batch normalization and a ReLU
activation, while the second produces the model predictions and thus has only
a single output channel. Since all elements of the target outputs were
constrained to be in the range [-1, 1], we apply a final tanh activation
function after this layer.

In all convolutional layers, the input is "same" padded to preserve height,
width, and depth dimensionality. Also, the convolutional layers in the
encoding and learning features phases are padded using torch’s circular
padding scheme to account for the periodic nature of the input Patterson maps.
Furthermore, all convolutional layers were initialized using the
kaiming_normal function of the default torch.nn module He et al. (2015). As
with the CrysFormer, our U-net model is robust to training batches of examples
with differing height, width, and depth shapes.

##  Appendix C Additional Details on Dataset Generation

To start preparing our dataset, we selected nearly 24000 representative
Protein Data Bank (PDB) entries using the following criteria: proteins solved
by X-ray crystallography after 1995, sequence length $\geq$ 40, refinement
resolution $\leq$ 2.75, refinement R-Free $\leq$ 0.28, with clustering at
30$\%$ sequence identity. The standardized modifications we applied to each
viable coordinate file were as follows: all temperature factors were set to
20, any selenomethionine residues were rebuilt as methionine, and all hydrogen
atoms were removed leaving only carbon, nitrogen, oxygen, and potentially
sulfur.

In our dataset generation process, an effort was taken to ensure diversity by
sampling from PDB entities with low sequence similarity to each other.
However, both test and training sets are taking random samples from the
conformations allowed in rotamer and Ramachandran space. Any similar
conformations would be expected to be in a different rotational orientation in
the cell by the nature of the selection process. We did not compute all-
versus-all clustering or force the test and training sets to sample distinct
conformational regions. For our 15-residue dataset, in order to obtain a
greater amount of starting coordinate files, we allowed at most 3 residues to
be shared between distinct examples. To prevent potential overfitting that
could arise from this sharing of subsegments, we enforced that all examples
derived from the same initial .pdb file would be placed together in either the
training or test set.

Another issue regarding ambiguity in Patterson map interpretation is the fact
that an electron density will always have the exact same Patterson map as its
corresponding centrosymmetry-related electron density. Hurwitz (2020) provided
a workaround that involved combining a set of atoms with the set of its
centrosymmetry-related atoms into a single example output. However, this also
requires a separate post-processing algorithm to separate the original and
centrosymmetric densities for each of his model’s predictions. Since we are
working with real-world structures –rather than randomly placed data– we can
exploit their known properties. In particular, we know that all proteinogenic
amino acids are naturally found in only one possible enantiomeric
configuration (Helmenstine, 2021). Although the mirror-image symmetry of
enantiomers is not exactly the same as centrosymmetry, we show that this is
enough to allow us to work with true electron densities of protein fragments.

##  Appendix D Description of Dataset Subset

Due to limitations of online storage space, we provide a subset of our
generated dataset. This subset represents a total of 200000 dipeptide
examples. As expected, patterson.tar.gz contains the generated Patterson maps,
while electron_density.tar.gz contains the corresponding electron densities.
Meanwhile, partial_structure.tar.gz contains both of the partial structures
for each dipeptide example in the subset.

The dataset can be downloaded through this link:

https://drive.google.com/drive/folders/1X7YkxDd7yTC1RTG1z3NbdRIfKLfFtkrx?usp=share_link

We will also provide a dataset of prepared .pdb coordinate files of 15-residue
examples, to which our dataset generation process can be applied in order to
produce Patterson map and electron density tensors.

##  Appendix E Additional Visualizations of Model Predictions

Figure LABEL:fig:AV represents an example in which the additional partial
structure input channels provided to the U-Net provided a substantial increase
in prediction quality, allowing it to produce a prediction similar to that of
the CrysFormer. Figure LABEL:fig:AL represents an example in which both
providing additional input channels to the U-Net and switching to CrysFormer
provided noticeable improvements in prediction quality.

It is clear that as prediction quality increases as indicated by reported
Pearson correlation, finer details of the true underlying structure are more
likely to be accurately reproduced. The predictions in Figure
LABEL:fig:more15mer (e), (f), (g), and (h), as well as Figure LABEL:pred15 (a)
[rank 55%] and (b) [rank 82%], were all successfully refined using all of the
mentioned autotracing and refinement procedures. But even for relatively poor
predictions such as (a) and (b), the rough overall shape can be reproduced
even though several portions have clear inaccuracies.

Figure 4: Scatterplot of the Pearson correlations of amplitudes of the poly-
alanine chains autotraced by shelxe to the ground truth amplitudes vs the
Pearson correlation of the predicted and ground truth maps for all 16,203 test
cases

Figure 4 shows the scatterplot of shelxe poly-alanine autotracing results on
the full 15-residue test set. As mentioned, examples for which the amplitudes
calculated from the initial poly-alanine chain built into the model electron
density prediction have a Pearson correlation coefficient with the true
underlying structure factor amplitudes of over 0.25 (shown above the red line)
are extremely likely to be successfully refined.


# Fairness-Driven Optimization of RIS-Augmented 5G Networks for Seamless 3D
UAV Connectivity Using DRL Algorithms

Yu Tian1, Ahmed Alhammadi1, Jiguang He1, Aymen Fakhreddine1, and Faouzi Bader1
1Technology Innovation Institute, Abu Dhabi, United Arab Emirates  E-mails:
{yu.tian, ahmed.alhammadi, jiguang.he, aymen.fakhreddine, carlos-
<EMAIL_ADDRESS>
###### Abstract

In this paper, we study the problem of joint active and passive beamforming
for reconfigurable intelligent surface (RIS)-assisted massive multiple-input
multiple-output systems towards the extension of the wireless cellular
coverage in 3D, where multiple RISs, each equipped with an array of passive
elements, are deployed to assist a base station (BS) to simultaneously serve
multiple unmanned aerial vehicles (UAVs) in the same time-frequency resource
of 5G wireless communications. With a focus on ensuring fairness among UAVs,
our objective is to maximize the minimum signal-to-interference-plus-noise
ratio (SINR) at UAVs by jointly optimizing the transmit beamforming parameters
at the BS and phase shift parameters at RISs. We propose two novel algorithms
to address this problem. The first algorithm aims to mitigate interference by
calculating the BS beamforming matrix through matrix inverse operations once
the phase shift parameters are determined. The second one is based on the
principle that one RIS element only serves one UAV and the phase shift
parameter of this RIS element is optimally designed to compensate the phase
offset caused by the propagation and fading. To obtain the optimal parameters,
we utilize one state-of-the-art reinforcement learning algorithm, deep
deterministic policy gradient, to solve these two optimization problems.
Simulation results are provided to illustrate the effectiveness of our
proposed solution and some insightful remarks are observed.

###### Index Terms:

Deep reinforcement learning, reconfigurable intelligent surfaces, joint active
and passive beamforming, UAV, interference mitigation.

##  I Introduction

Thanks to their convenient deployment, simple structure, and versatile landing
capabilities, drones, also known as unmanned aerial vehicles (UAVs), have
recently experienced significant growth. This expansion has given rise to a
diverse array of applications, including but not limited to goods delivery,
urban air taxis, remote surveillance, border control, agricultural and
industrial monitoring, as well as disaster relief [1, 2]. Despite the diverse
nature of these applications, they share a common critical requirement: the
necessity for robust three-dimensional (3D) wireless connectivity to
facilitate the seamless exchange of real-time sensor data and control
commands. The paramount concern across these applications is to ensure that
the connectivity is both reliable and secure. This is imperative to prevent
scenarios where prolonged control loss could potentially result in severe
accidents.

Current widely deployed commercial fifth-generation (5G) base stations (BSs)
are a suitable option for providing connectivity to UAVs. However, these BSs
are specifically designed to offer two-dimensional (2D) coverage for ground
users, resulting in weakened signal strength in the sky. Addressing the need
for ubiquitous 3D coverage and global connectivity is a key focus for future
wireless networks, i.e., beyond 5G and sixth generation (6G) [3]. To overcome
the absence of dedicated infrastructure for airborne UAVs, solutions that
minimize the need for significant hardware upgrades or additional BS
deployments by mobile network operators are preferable. Emphasizing reliance
on existing 5G cellular networks and their enhancement to meet quality-of-
service (QoS) requirements for UAV communications is essential. The authors of
[4] explore the feasibility of utilizing reconfigurable intelligent surfaces
(RISs) to extend 3D coverage in cellular networks for aerial users. Their
study indicates that deploying RISs can yield impressive gains in terms of
achievable rate, outperforming traditional RIS-free 5G networks. However, the
optimal power allocation and phase shift configuration for multi-RIS-augmented
3D coverage in multi-user cellular networks remain unexplored.

In the realm of recent research on RIS-augmented wireless communication
systems, the predominant focus has been on two primary objectives for power
allocation and RIS configuration optimization: maximizing the sum data rate
[5, 6, 7, 8, 9] and maximizing the minimum signal-to-interference-plus-noise
ratio (SINR) [10, 11]. The former objective centers around maximizing the
overall throughput of the system, prioritizing system-wide efficiency. In
contrast, the latter emphasizes the maximization of user fairness, with a
focus on equitable resource distribution among users. As highlighted in [12],
the existing literature predominantly employs deep reinforcement learning
(DRL) to address throughput optimization challenges, while there is a
noticeable scarcity of studies utilizing this approach to specifically tackle
user fairness concerns. Capitalizing on the demonstrated success of DRL in
addressing NP-hard non-convex optimization problems, our paper delves into the
user fairness within a 3D multi-RIS-augmented multiple-input multiple-output
(MIMO) system. Our objective is to maximize the minimum SINR through the joint
optimization of transmit beamforming vectors and phase shift parameters. To
achieve this and alleviate the complexities associated with the massive search
space, we introduce the innovative “One RIS Element Serving One UAV” (ORESOU)
approach and the interference canceling (IC) method. Leveraging cutting-edge
DRL algorithm, deep deterministic policy gradient (DDPG), we implement and
evaluate the performance of our proposed algorithms.

##  II System Model

The configuration of the system model is depicted in Fig. 1. Our focus is on a
massive MIMO system aided by multiple RISs. The system comprises a single base
station (BS) equipped with $N_{\text{B}}$ vertical antenna elements,
$N_{\text{R}}$ RISs, and $N_{\text{U}}$ UAVs, each equipped with a single
omnidirectional antenna. Each RIS consists of $N=N_{x}\times N_{y}$ reflecting
meta-atoms, where $N_{x}$ and $N_{y}$ denote the number of RIS meta-atoms
along the $x$ and $y$ axes, respectively. For the sake of simplifying our
analytical approach, we make the assumption that the RIS is square, possessing
an equal number of meta-atoms along both the $x$ and $y$ axes, i.e.,
$N_{x}=N_{y}$.

Figure 1: 3D connectivity enabled by multiple RISs along with a legacy 5G BS.

We assume that the BS exclusively serves terrestrial users, and its beams are
down-tilted. Simultaneously, the UAVs operate at a substantial altitude,
beyond the terrestrial coverage area. Consequently, direct links between the
BS and UAVs are nonexistent. In this context, RISs can be strategically
positioned on the ground, walls, or atop buildings to enhance communication
between the BS and UAVs by intelligently reflecting signals from the BS. All
the RISs are connected and receive a feedback channel from the BS.

Let $\mathbf{G}_{\text{B},n_{r}}\in\mathbb{C}^{{N_{\text{B}}}\times N}$ denote
the channel from the BS to the $n_{r}$-th RIS, and
$\mathbf{h}_{n_{r},n_{u}}\in\mathbb{C}^{1\times N}$ represent the channel from
the $n_{r}$-th RIS to the $n_{u}$-th UAV. Typically, RISs are positioned in
locations with a direct link to the BS to ensure high signal quality.
Nevertheless, the presence of nearby objects inevitably leads to pervasive
reflection, inducing small-scale fading. Therefore, we model the BS-RIS
channel by employing the Rician fading which is a summation of weighted line-
of-sight (LoS) and non-line-of-sight (NLoS) paths, expressed mathematically
as:

$\displaystyle\mathbf{G}_{\text{B},n_{r}}=\underbrace{\sqrt{\frac{1}{\kappa_{1}+1}}\widetilde{\mathbf{G}}_{n_{r}}}_{\text{NLoS}}+\underbrace{\sqrt{\frac{\kappa_{1}}{\kappa_{1}+1}}\overline{\mathbf{G}}_{n_{r}}}_{\text{LoS}},$
(1)

where $\widetilde{\mathbf{G}}_{n_{r}}$ represents the NLoS component,
$\overline{\mathbf{G}}_{n_{r}}$ is the LoS component, and $\kappa_{1}$ denotes
the power ratio between $\overline{\mathbf{G}}_{n_{r}}$ and
$\widetilde{\mathbf{G}}_{n_{r}}$. Each entry of
$\widetilde{\mathbf{G}}_{n_{r}}$ follows a complex Gaussian distribution
$\mathcal{CN}(0,1)$. $\overline{\mathbf{G}}_{n_{r}}$ is expressed with respect
to the steering vector for the square RIS with ideal isotropic elements and
the steering vector of the BS. These vectors depend on the azimuth and
elevation angles of arrival and departure for the signals. Specifically,

$\displaystyle\overline{\mathbf{G}}_{n_{r}}=\boldsymbol{\alpha}_{\text{R}}(\theta^{r}_{\text{R},n_{r}},\phi^{r}_{\text{R},n_{r}})\boldsymbol{\alpha}_{\text{B}}^{\mathsf{H}}(\phi^{t}_{\text{B},n_{r}}-\beta)\exp{(-j2\pi
f_{c}\tau_{\text{B},n_{r}})},$ (2)

where $\theta^{r}_{\text{R},n_{r}}\in[0,\pi]$,
$\phi^{r}_{\text{R},n_{r}}\in[-\pi/2,\pi/2]$, and
$\phi^{t}_{\text{B},n_{r}}\in[-\pi/2,\pi/2]$ represent the azimuth, elevation
angle of arrival (AoA), and elevation angle of departure (AoD), respectively,
$\beta$ denotes the down-tilted orientation at the BS. Here, we assume that
all the communications share the same time-frequency resources, $f_{c}$
denotes the carrier frequency, and $\tau_{\text{B},n_{r}}$ is the propagation
delay from the BS to the $n_{r}$-th RIS.
$\boldsymbol{\alpha}_{\text{R}}(\theta^{r}_{\text{R},n_{r}},\phi^{r}_{\text{R},n_{r}})$
and $\boldsymbol{\alpha}_{\text{B}}(\phi^{t}_{\text{B},n_{r}})$ are the
normalized array response steering vectors, defined as

$\displaystyle\boldsymbol{\alpha}_{\text{R}}(\theta^{r}_{\text{R},n_{r}},\phi^{r}_{\text{R},n_{r}})$
$\displaystyle\triangleq\frac{1}{\sqrt{N}}\Bigg{[}1,\exp\left(j\frac{2\pi
d_{x}}{\lambda}\cos(\theta^{r}_{\text{R},n_{r}})\sin(\phi^{r}_{\text{R},n_{r}})\right),$
$\displaystyle\;\;\;\cdots,\exp\left(j\frac{2\pi
d_{x}}{\lambda}(N_{x}-1)\cos(\theta^{r}_{\text{R},n_{r}})\sin(\phi^{r}_{\text{R},n_{r}})\right)\Bigg{]}^{\mathsf{T}}$
$\displaystyle\otimes\Bigg{[}1,\exp\left(j\frac{2\pi
d_{y}}{\lambda}\sin(\theta^{r}_{\text{R},n_{r}})\sin(\phi^{r}_{\text{R},n_{r}})\right),$
$\displaystyle\;\;\;\cdots,\exp\left(j\frac{2\pi
d_{y}}{\lambda}(N_{y}-1)\sin(\theta^{r}_{\text{R},n_{r}})\sin(\phi^{r}_{\text{R},n_{r}})\right)\Bigg{]}^{\mathsf{T}},$
(3)
$\displaystyle\boldsymbol{\alpha}_{\text{B}}(\phi^{t}_{\text{B},n_{r}}-\beta)\triangleq\frac{1}{\sqrt{N_{\text{B}}}}\Bigg{[}1,\exp\left(j\frac{2\pi
d_{z}}{\lambda}\cos(\phi^{t}_{\text{B},n_{r}}-\beta)\right),$
$\displaystyle\;\;\;\cdots,\exp\left(j\frac{2\pi
d_{z}}{\lambda}(N_{\text{B}}-1)\cos(\phi^{t}_{\text{B},n_{r}}-\beta)\right)\Bigg{]}^{\mathsf{T}},$
(4)

where $d_{x}$, $d_{y}$, and $d_{z}$ represent the inter-element spacing across
$x$, $y$, and $z$ axes, respectively. $\otimes$ denotes the Kronecker product.
We can easily see that
$\|\boldsymbol{\alpha}_{\text{R}}(\theta^{r}_{\text{R},n_{r}},\phi^{r}_{\text{R},n_{r}})\|_{2}=\|\boldsymbol{\alpha}_{\text{B}}(\phi^{t}_{\text{B},n_{r}}-\beta)\|_{2}=1$.

Meanwhile, the channels between the $n_{r}$-th RISs and the $n_{u}$-th UAVs
are also modeled by Rician fading, expressed as

$\displaystyle\mathbf{h}_{n_{r},n_{u}}=\sqrt{\frac{\kappa_{2}}{\kappa_{2}+1}}\overline{\mathbf{h}}_{n_{r},n_{u}}+\sqrt{\frac{1}{\kappa_{2}+1}}\widetilde{\mathbf{h}}_{n_{r},n_{u}},$
(5)

where each element of $\widetilde{\mathbf{h}}_{n_{r},n_{u}}$ follows
$\mathcal{CN}(0,1)$ and
$\overline{\mathbf{h}}_{n_{r},n_{u}}=\boldsymbol{\alpha}_{\text{R}}^{\mathsf{H}}(\theta^{t}_{n_{r},n_{u}},\phi^{t}_{n_{r},n_{u}})\times\exp{(-j2\pi
f_{c}\tau_{n_{r},n_{u}})}$ with $\theta^{t}_{n_{r},n_{u}}$ and
$\phi^{t}_{n_{r},n_{u}}$ denoting the azimuth and elevation AoDs from the RIS,
and $\tau_{n_{r},n_{u}}$ representing the propagation delay between the
$n_{r}$-th RIS and the $n_{u}$-th UAV.

Hence, the channel from the BS to the $n_{u}$-th UAV through the
$N_{\text{R}}$ RISs can be expressed as

$\displaystyle\boldsymbol{\omega}_{\text{B},n_{u}}=\sum_{n_{r}=1}^{N_{\text{R}}}\mathbf{h}_{n_{r},n_{u}}\mathbf{\Phi}_{n_{r}}\mathbf{G}_{\text{B},n_{r}}$
(6)

where

$\displaystyle\mathbf{\Phi}_{n_{r}}=\mathrm{diag}(\exp{(\theta_{n_{r},1})},...,\exp{(\theta_{n_{r},N})})$
(7)

is the diagonal phase control matrix of the $n_{r}$-th RIS and
$\mathrm{diag}(\mathbf{a})$ denotes a diagonal matrix with the entries of
$\mathbf{a}$ on its diagonal.

Furthermore, the received signal at the $n_{u}$-th UAV comprises the desired
signal, interference signal from the other UAVs, and the noise, expressed as

$\displaystyle
y_{n_{u}}=\boldsymbol{\omega}_{\text{B},n_{u}}\left(\sqrt{P_{n_{u}}}\mathbf{f}_{n_{u}}s_{n_{u}}+\sum_{i=1,i\neq
n_{u}}^{N_{\text{U}}}\sqrt{P_{i}}\mathbf{f}_{i}s_{i}\right)+z_{n_{u}},$ (8)

where $P_{n_{u}}$, $\mathbf{f}_{n_{u}}\in\mathbb{C}^{N_{\text{B}}}$, and
$s_{n_{u}}\sim\mathcal{CN}(0,1)$ are respectively the transmit power, the
beamforming vector, and the transmitted signal symbol for the ${n_{u}}$-th
UAV. $z_{n_{u}}$ is the additive complex Gaussian noise with zero mean and
variance of $\sigma_{z}^{2}$, i.e., $\mathcal{CN}(0,\sigma_{z}^{2})$.

Finally, the SINR of the $n_{u}$-th UAV is expressed as

$\displaystyle\gamma_{n_{u}}=\frac{P_{n_{u}}|\boldsymbol{\omega}_{\text{B},n_{u}}\mathbf{f}_{n_{u}}|^{2}}{\sum_{i\neq
n_{u}}P_{i}|\boldsymbol{\omega}_{\text{B},n_{u}}\mathbf{f}_{i}|^{2}+\sigma_{z}^{2}},$
(9)

and the sum rate of all the $N_{\text{U}}$ UAVs is presented as

$\displaystyle R=\sum_{n_{u}=1}^{N_{\text{U}}}\log_{2}(1+\gamma_{n_{u}}).$
(10)

##  III Problem Formulation

Most of the existing literature [5, 6, 7, 8, 9] focus on maximizing the sum
rate of all users, potentially leading to fairness concerns as illustrated in
Fig. 2(a). This objective function tends to allocate more power to users with
superior channel gains while less power to those with inferior channel gains.
Motivated by [10], we opt for the objective function of maximizing the minimal
SINR. The optimization problem can be expressed as

$\displaystyle\max\limits_{\mathbf{\Phi},\mathbf{F},\mathbf{P}}\,\,\min\limits_{n_{u}}\gamma_{n_{u}}$
(11) $\displaystyle\text{s.t.}\begin{cases}||\mathbf{P}||_{1}\leq P_{\max}\\\
\mathrm{diag}(\mathbf{F}^{\mathsf{H}}\mathbf{F})=\mathbf{I}_{N_{\text{U}}}\\\
|\mathbf{\Phi}|=\mathbf{I}_{NN_{\text{R}}}\\\ \end{cases},$ (12)

where
$\mathbf{P}=\mathrm{diag}(P_{1},...,P_{N_{\text{U}}})\in\mathbb{R}^{N_{\text{U}}\times
N_{\text{U}}}$,
$\mathbf{F}=[\mathbf{f}_{1},...,\mathbf{f}_{N_{\text{U}}}]\in\mathbb{C}^{N_{\text{B}}\times
N_{\text{U}}}$,
$\mathbf{\Phi}=\mathrm{diag}(\mathbf{\Phi}_{1},...,\mathbf{\Phi}_{N_{\text{R}}})\in\mathbb{C}^{NN_{\text{R}}\times
NN_{\text{R}}}$, and $P_{\max}$ represents the maximum transmit power at the
BS.

The optimization problem (11) is non-convex, involving three variable
matrices. Presently, many researchers resort to reinforcement learning (RL) to
jointly optimize all the three variables simultaneously, which suffers from a
vast search space. In this paper, we propose two algorithms to tackle this
issue.

###  III-A Interference Cancellation (IC) Algorithm

The overarching objective of maximizing the minimum SINR in (11) is to
simultaneously maximize the power of the received desired signal and minimize
the interference power. The ideal outcome is characterized by zero
interference power, with all users sharing the same maximized received power.
Consequently, the objective function can be formulated as:

$\displaystyle\max\limits_{\mathbf{\Phi},\mathbf{F},\mathbf{P}}\,\,\alpha$
(13)
$\displaystyle\text{s.t.}\begin{cases}\mathbf{H}\mathbf{\Phi}\mathbf{G}\mathbf{F}\mathbf{P}=\alpha\mathbf{I}_{N_{\text{U}}}\\\
||\mathbf{P}||_{1}\leq P_{\max}\\\
\mathrm{diag}(\mathbf{F}^{\mathsf{H}}\mathbf{F})=\mathbf{I}_{N_{\text{U}}}\\\
|\mathbf{\Phi}|=\mathbf{I}_{NN_{\text{R}}}\\\ \end{cases}$ (14)

where
$\mathbf{G}=[\mathbf{G}_{\text{B},1},...,\mathbf{G}_{\text{B},N_{\text{R}}}]\in\mathbb{C}^{NN_{\text{R}}\times
N_{\text{B}}}$ and
$\mathbf{H}=[[\mathbf{h}_{1,1},...,\mathbf{h}_{NN_{\text{R}},1}];...;[\mathbf{h}_{1,N_{\text{U}}},...,\mathbf{h}_{NN_{\text{R}},N_{\text{U}}}]]\in\mathbb{C}^{N_{\text{U}}\times
NN_{\text{R}}}$.

Because $\mathbf{\Phi}$ and $\mathbf{P}$ are diagonal,
$\mathbf{H}\mathbf{\Phi}\mathbf{G}\in\mathbb{C}^{N_{\text{U}}\times
N_{\text{B}}}$ and $\mathbf{F}\mathbf{P}\in\mathbb{C}^{N_{\text{B}}\times
N_{\text{U}}}$ should be reciprocal to meet the first condition in (12).
$\mathbf{FP}$ can be combined together to one variable
$\mathbf{\hat{F}}=\mathbf{FP}\in\mathbb{C}^{N_{\text{B}}\times N_{\text{U}}}$
with condition
$\mathrm{trace}(\mathbf{\hat{F}}^{\mathsf{H}}\mathbf{\hat{F}})\leq P_{\max}$.
Once $\mathbf{\Phi}$ is obtained, $\mathbf{\hat{F}}$ can be obtained by the
pseudo inverse of $\mathbf{H}\mathbf{\Phi}\mathbf{G}$,
$\mathbf{\hat{F}}=\hat{P}\frac{\mathbf{\widetilde{F}}}{\mathrm{trace}(\mathbf{\widetilde{F}}^{\mathsf{H}}\mathbf{\widetilde{F}})}$
with $\mathbf{\widetilde{F}}=\mathrm{pinv}(\mathbf{H}\mathbf{\Phi}\mathbf{G})$
and $\hat{P}\leq P_{\max}$. In this way, we can get

$\displaystyle\mathbf{H}\mathbf{\Phi}\mathbf{G}\mathbf{\hat{F}}=\hat{P}\frac{\mathbf{H}\mathbf{\Phi}\mathbf{G}\mathbf{\widetilde{F}}}{\mathrm{trace}(\mathbf{\widetilde{F}}^{\mathsf{H}}\mathbf{\widetilde{F}})}=\hat{P}\frac{\mathbf{I}_{N_{\text{U}}}}{\mathrm{trace}(\mathbf{\widetilde{F}}^{\mathsf{H}}\mathbf{\widetilde{F}})}.$
(15)

Note that $\mathbf{H}\mathbf{\Phi}\mathbf{G}$ and $\mathbf{\hat{F}}$ are non-
square matrices. They should be full-rank matrices to satisfy the pseudo
inverse requirements. Namely,
$\mathrm{rank}(\mathbf{H}\mathbf{\Phi}\mathbf{G})=\mathrm{rank}(\mathbf{\hat{F}})=\min{(N_{\text{U}},N_{\text{B}})}$.

To maximize above equation, $\hat{P}$ should be equal to $P_{\max}$, that is,
$\hat{P}=P_{\max}$. Finally, we can get

$\displaystyle\mathbf{\hat{F}}=P_{\max}\frac{\mathbf{\widetilde{F}}}{\mathrm{trace}(\mathbf{\widetilde{F}}^{\mathsf{H}}\mathbf{\widetilde{F}})}.$
(16)

Because the interference is canceled, the SINR of the $n_{u}$-th UAV in (9)
becomes SNR, that is,
$\gamma_{n_{u}}={P_{n_{u}}|\boldsymbol{\omega}_{\text{B},n_{u}}\mathbf{f}_{n_{u}}|^{2}}/{\sigma_{z}^{2}}$.
In this context, noise power plays a vital role in minimum SINR optimization.

###  III-B One-RIS-Element-Serving-One-UAV (ORESOU) Algorithm

$\mathbf{H}\mathbf{\Phi}\mathbf{G}$ and $\mathbf{\hat{F}}$ may not always
satisfy the full-rank requirement. To address this issue, we propose the
ORESOU algorithm with an assumption that each RIS element serves only one UAV.
The phase shift of this element is adjusted to compensate for the phase offset
induced by the propagation and fading from the BS to the RIS element and from
the RIS element to the target UAV. In this context, we can substitute the
continuous $\mathbf{\Phi}$ with an association matrix
$\mathbf{\overline{\Phi}}\in\mathbb{B}^{NN_{\text{R}}\times N_{\text{U}}}$
that links all the RIS elements and all the UAVs. $\mathbf{\overline{\Phi}}$
possesses the property that its elements take only two values from
$\mathbb{B}=\\{0,1\\}$ and the sum of each row equals 1.

##  IV Deep Reinforcement Learning Framework

In this section, we provide a brief overview of DRL, laying the foundation for
the proposed joint design of BS transmit beamforming and RIS phase shifts.

###  IV-A Overview of RL

RL is a branch of machine learning that deals with how agents should take
actions in an environment to maximize some notion of cumulative reward [5]. In
a typical RL, the agent incrementally refines its decision-making capabilities
through ongoing interactions with the environment. During each interaction,
the agent applies actions to the environment, observing the instant rewards
and tracking state transitions of the environment.

The primary objective of RL is to maximize the cumulative sum of rewards, as
quantified by the state-action value functions. Given a state, RL can
determine the optimal action based on the trained policy. In traditional RL
algorithms, the state-action value function is computed and updated using the
Bellman equation. To manage discrete states and actions, all corresponding
values are stored in a Q-table. However, when dealing with continuous actions,
such as beamforming and phase shift matrices in our problem, the Q-table
becomes impractical. In such cases, a deep neural network is employed as a
replacement for the Q-table, allowing for a more flexible and scalable
representation of the state-action space which facilitates the creation of the
DRL framework.

###  IV-B Actor-Critic Framework

The RL actor-critic framework is a hybrid architecture that combines a policy-
based (actor) and a value-based (critic) approach. Actor defines the agent’s
policy that determines the probability distribution of actions for a given
state, while critic evaluates the value of these actions by estimating a
state-action value function. Actor evaluates the value of these actions based
on the information from the critic’s dominance information to update it,
guiding the selection of actions that lead to higher than expected rewards. At
the same time, critic is trained to minimise the time-difference error between
the estimated and actual rewards, providing a benchmark for the actor’s
updates. This cooperative learning strategy improves stability and efficiency
compared to purely policy-based or value-based approaches. The actor-critic
framework, typically implemented using deep neural networks, is suitable for
tasks with a continuous action space and has been successful in a variety of
applications, including robot control and game play. Actor-critic frameworks
implemented via deep neural networks (DNNs) perform well in continuous action
space scenarios, overcoming the limitations of traditional Q-table approaches.

###  IV-C Deep Deterministic Policy Gradient (DDPG)

DDPG is a prominent off-policy actor-critic algorithm designed to solve RL
problems in continuous action spaces. DDPG is characterised by the learning of
deterministic policies, which allows for direct mapping from states to
actions, which is particularly beneficial in scenarios with continuous and
high-dimensional action spaces [13]. The algorithm maintains a network of
actors, which are responsible for policy approximation, and a network of
critics, which evaluate the expected payoffs based on the actions chosen by
the actors. DDPG employs experience replay, a mechanism for storing and
randomly sampling past experience, which enhances the stability and efficiency
of the learning by breaking down correlations between consecutive samples. In
addition, the integration of a target network that periodically updates slower
copies of the main network helps to achieve a more stable training process by
providing consistent target values during learning updates.

The DDPG process involves iteratively updating actor and critic networks
through a combination of policy gradient ascent and temporal difference
learning. The actor is trained to maximise the expected return from the
critic’s perspective, and the critic is updated to minimise the mean-squared
time-difference error between its predicted Q-value and the target Q-value
obtained from the target network. This process is guided by the deterministic
policy gradient theorem, which drives the actor’s updates in the direction of
maximising the expected return. The algorithm also explores by adding noise to
the selection of actions, prompting a balance between exploration and
exploitation. Overall, DDPG provides an efficient and stable solution to the
continuous control problem and works in line with the principles of strategy
optimisation and value estimation in the actor-critic framework.

Algorithm 1 DDPG for IC and ORESOU algorithms

Input: $\mathbf{G}$, $\mathbf{H}$
Output: Q-value function, $\mathbf{\Phi}$ for IC,
$\\{\mathbf{\overline{\Phi}},\hat{\mathbf{F}}\\}$ for ORESOU Initialization:
Critic, Actor, and target networks, replay buffer

1:for each episode do

2: Initialize $\mathbf{G}$ and $\mathbf{H}$, and receive first state

3: for each time step do

4: Obtain action ($\mathbf{\Phi}$ or
$\\{\mathbf{\overline{\Phi}},\hat{\mathbf{F}}\\}$) from Actor

5: Observe new state given action and noise, calculate best $\mathbf{\hat{F}}$
for IC and best $\mathbf{\Phi}$ for ORESOU

6: Obtain minimum SINR

7: Store state, action, new state, and minimum SINR to replay buffer

8: Obtain Q-value from Critic

9: Sample random mini-batches from replay buffer

10: Update Critic by temporal difference learning

11: Update Actor by policy gradient ascent

12: Periodically slowly update target networks

13: end for

14:end for

###  IV-D DDPG for IC and ORESOU

The implementation details of DDPG for IC and ORESOU are described in
Algorithm 1. Both the IC and ORESOU adhere to the same reward and
environmental setup. The reward is defined as the minimum SINR, emphasizing
the optimization objective. The environmental setup encompasses the spatial
coordinates and channel information associated with the BS, RISs, and UAVs.

The primary distinction lies in the action space. As outlined in Section
III-A, IC leverages DDPG to explore and determine $\mathbf{\Phi}$,
subsequently obtaining $\mathbf{\hat{F}}$ through equation (16). In contrast,
ORESOU employs DDPG to explore both $\mathbf{\overline{\Phi}}$ and
$\mathbf{\hat{F}}$. Consequently, the action space for IC encompasses only
$\mathbf{\Phi}$, while ORESOU’s action space encompasses both
$\mathbf{\overline{\Phi}}$ and $\mathbf{\hat{F}}$. This differentiation
captures the key distinction in the action formulation between the two
algorithms. In comparison to conventional DDPG, where the action space
encompasses $\mathbf{\Phi}$ and $\mathbf{\hat{F}}$, IC simplifies the action
space by disregarding $\mathbf{\hat{F}}$. On the other hand, ORESOU
streamlines the action space by employing binary $\mathbf{\overline{\Phi}}$ in
place of the continuous $\mathbf{\Phi}$.

##  V Numerical Simulations

In this section, we conduct a comprehensive performance evaluation of IC and
ORESOU algorithms, both implemented using the DDPG algorithm. Simultaneously,
to highlight the advantages of our proposed methodologies, we compare them
against two baseline approaches: a conventional DRL method and a random search
strategy.

The DRL method involves utilizing DDPG to explore the continuous phase shift
configuration $\mathbf{\Phi}$ and the BS beamforming matrix $\mathbf{F}$.
Conversely, the random search approach generates $\mathbf{\Phi}$ and
$\mathbf{F}$ randomly, ensuring adherence to the conditions specified in
equation (14). For the DDPG implementation, we adopt the well-established
implementation available on GitHub [14]. Key system setting parameters are
detailed in Table I, and we focus on a simplified scenario featuring 2 RISs
and 2 UAVs.

TABLE I: PARAMETER VALUES USED FOR THE SYSTEM MODEL Parameter | Value
---|---
BS coordinates (m) | (0,0,2)
RIS coordinates1 (m) | (-2.5,8,0),(2.5,0,0)
UAV coordinates1 (m) | (-3,10,6),(2,6,10)
$\beta$ | 0
$N_{\text{B}}$ | 4
$N_{\text{U}}$ | 2
$N_{\text{R}}$ | 2
$N$ | 16
$f_{c}$ | 28 GHz
$\kappa_{1}$,$\kappa_{2}$ | 30 dB
$\sigma_{z}^{2}$ | -100 dBm

(a) Objective: Maximize the sum rate.

(b) Objective: Maximize the minimum SINR.

Figure 2: Instant and average SINRs of two UAVs using two objectives (rolling average window = 300 time steps). Figure 3: Minimum SINR versus $P_{\max}$. TABLE II: Comparison between one-bit and continuous RIS with $P_{\max}=45$ dBm Minimum SINR (dB) | IC | ORESOU | DRL | Random Search
---|---|---|---|---
One-bit | 11.63 | 10.167 | 9.79 | 4.058
Continuous | 15.36 | 12.05 | 11.52 | 10.10
Difference | 3.73 | 1.89 | 1.73 | 6.05

We first investigate the issue caused by setting the objective function as
maximizing the sum rate. In this experiment, we use DDPG to explore
$\mathbf{\Phi}$ and $\mathbf{F}$ to achieve the maximum sum rate with
$P_{\max}=35$ dBm. The recorded instant and average SINRs of two UAVs at each
time step are shown in Fig. 2(a). Note that the average SINR is computed by a
rolling window of 300 time steps. We can see that the SINR of ‘UAV 1’ tends to
very low while the SINR of ‘UAV 2’ becomes very high. In reference to Table I,
we can observe that ‘UAV 2’ experiences less propagation loss from the 2 RISs
compared to ‘UAV 1’. Consequently, it can be deduced that ‘UAV 2’ possesses a
superior channel condition. We can conclude that DRL tends to assign most of
the resources to the UAV with the best channel condition to achieve the
largest SINR and the maximum sum rate. When formulating the objective function
to maximize the minimum SINR, the issue of SINR imbalance between the two UAVs
is negligible, as illustrated in Fig. 2(b). The discrepancy in both instant
and average SINRs for these two UAVs is much smaller, emphasizing the
consideration for fairness among different UAVs. Therefore, we opt for
maximizing the minimum SINR as our objective function.

In the evaluation of the IC, ORESOU, DRL, and random search algorithms, we
analyze their best minimum SINRs for ‘UAV 1’ and ‘UAV 2’ over 10,000
iterations with a varying $P_{\max}$ ranging from 20 to 45 dBm, as depicted in
Fig. 3. Notably, we observe an increasing trend in minimum SINR with the
increment of $P_{\max}$. Among the algorithms, IC consistently outperforms the
others, exhibiting superior performance in terms of minimum SINR. ORESOU
closely follows IC in performance, demonstrating a slightly lower but
significantly improved performance compared to the two baseline approaches,
DRL, and random search. In comparison to DRL, IC demonstrates an enhancement
ranging from 1.5 dB (for $P_{\max}=20$ dBm) to 4 dB ($P_{\max}=45$ dBm), while
ORESOU exhibits improvement spanning from 0.5 dB ($P_{\max}=25$ dBm) to 3.3 dB
($P_{\max}=40$ dBm).

In the context of practical engineering implementation, the complexity
associated with continuous phase configuration in RIS systems has led
researchers to explore the viability of employing a simplified one-bit phase
configuration. In this approach, the phase shift of each RIS element can only
take on two values: 0 or $\pi$, significantly reducing the implementation
complexity. In our evaluation, we compared the performance of our proposed
algorithms and baseline approaches under the one-bit RIS configuration with a
maximum power constraint of $P_{\max}=45$ dBm. The results, as presented in
Table II, reveal that adopting a one-bit RIS configuration can lead to
substantial performance degradation across all four algorithms. Specifically,
the Minimum SINR of IC experiences a significant 3.73 dB reduction. Both
ORESOU and DRL algorithms show a similar and moderate reduction of about 1.8
dB. In contrast, the random search approach exhibits the most substantial
performance degradation, with a fall of over 6 dB. Despite these challenges,
IC and ORESOU still outperform the DRL and random search approaches in this
one-bit RIS setting. The consistent superiority of IC and the competitive
performance of ORESOU underscore the efficacy of our proposed algorithms in
optimizing resource allocation and tackling interference mitigation issues.

##  VI Conclusions

In this paper, we have focused on fostering fairness among various UAVs within
a multi-RIS-augmented 5G MIMO network, aiming for seamless 3D connectivity.
Our specific goal is to maximize the minimum SINR by concurrently optimizing
transmit beamforming vectors and RIS phase shift parameters. To address the
intricate NP-hard non-convex problem, we have employed the cutting-edge DDPG
algorithm. Recognizing the challenge posed by the extensive action search
space inherent to DRL, we have introduced two novel algorithms, ORESOU and IC.
Through comprehensive numerical simulations, we have thoroughly analyzed and
compared the performance of ORESOU, IC, DRL, and random search algorithms. Our
results have highlighted that IC and ORESOU not only effectively reduce the
action space but also significantly enhance performance, particularly in
achieving the best minimum SINR. These findings verified the efficacy of our
proposed algorithms in optimizing the performance of multi-RIS-augmented 5G
MIMO networks.

## References

  * [1] C. Zhang, W. Zhang, W. Wang, L. Yang, and W. Zhang, “Research challenges and opportunities of UAV millimeter-wave communications,” _IEEE Wireless Commun._ , vol. 26, no. 1, pp. 58–62, 2019.
  * [2] Y. Tian, G. Pan, M. A. Kishk, and M.-S. Alouini, “Stochastic analysis of cooperative satellite-UAV communications,” _IEEE Trans. Wireless Commun._ , vol. 21, no. 6, pp. 3570–3586, 2022.
  * [3] Z. Zhang, Y. Xiao, Z. Ma, M. Xiao, Z. Ding, X. Lei, G. K. Karagiannidis, and P. Fan, “6G wireless networks: Vision, requirements, architecture, and key technologies,” _IEEE Veh. Technol. Mag._ , vol. 14, no. 3, pp. 28–41, 2019.
  * [4] J. He, A. Fakhreddine, A. S. de Sena, Y. Tian, and M. Debbah, “Unleashing 3D connectivity in beyond 5G networks with reconfigurable intelligent surfaces,” _arXiv preprint arXiv:2305.04661_ , 2023.
  * [5] C. Huang, R. Mo, and C. Yuen, “Reconfigurable intelligent surface assisted multiuser MISO systems exploiting deep reinforcement learning,” _IEEE J. Sel. Areas Commun._ , vol. 38, no. 8, pp. 1839–1850, 2020.
  * [6] B. Saglam, D. Gurgunoglu, and S. S. Kozat, “Deep reinforcement learning based joint downlink beamforming and RIS configuration in RIS-Aided MU-MISO systems under hardware impairments and imperfect CSI,” in _2023 IEEE Int. Conf. Commun. Workshops (ICC Workshops)_ , 2023, pp. 66–72.
  * [7] H. Chen, G. Yang, and Y.-C. Liang, “Joint active and passive beamforming for reconfigurable intelligent surface enhanced symbiotic radio system,” _IEEE Wireless Commun. Lett._ , vol. 10, no. 5, pp. 1056–1060, 2021.
  * [8] Y. Xu, Z. Gao, Z. Wang, C. Huang, Z. Yang, and C. Yuen, “RIS-enhanced WPCNs: Joint radio resource allocation and passive beamforming optimization,” _IEEE Trans. Veh. Technol._ , vol. 70, no. 8, pp. 7980–7991, 2021.
  * [9] K. Stylianopoulos, G. Alexandropoulos, C. Huang, C. Yuen, M. Bennis, and M. Debbah, “Deep contextual bandits for orchestrating multi-user MISO systems with multiple RISs,” in _ICC 2022 - IEEE Int. Conf. Commun._ , 2022, pp. 1556–1561.
  * [10] X. Li, J. Fang, F. Gao, and H. Li, “Joint active and passive beamforming for intelligent reflecting surface-assisted massive mimo systems,” _arXiv preprint arXiv:1912.00728_ , 2019.
  * [11] H. Lu, Y. Zeng, S. Jin, and R. Zhang, “Aerial intelligent reflecting surface: Joint placement and passive beamforming design with 3D beam flattening,” _IEEE Trans. Wireless Commun._ , vol. 20, no. 7, pp. 4128–4143, 2021.
  * [12] H. Zhou, M. Erol-Kantarci, Y. Liu, and H. V. Poor, “A survey on model-based, heuristic, and machine learning optimization approaches in RIS-aided wireless networks,” _arXiv preprint arXiv:2303.14320_ , 2023.
  * [13] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine, “Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor,” in _Int. Conf. Mach. Learn._ PMLR, 2018, pp. 1861–1870.
  * [14] B. Saglam, “RIS MISO deep reinforcement learning,” https://github.com/baturaysaglam/RIS-MISO-Deep-Reinforcement-Learning, 2021.