# What is a meaningful representation of protein sequences?

Nicki Skafte Detlefsen
Section for Cognitive Systems
Technical University of Denmark
<EMAIL_ADDRESS>
Søren Hauberg
Section for Cognitive Systems
Technical University of Denmark
<EMAIL_ADDRESS>Wouter Boomsma
Department of Computer Science
University of Copenhagen
<EMAIL_ADDRESS>

###### Abstract

How we choose to represent our data has a fundamental impact on our ability to
subsequently extract information from them. Machine learning promises to
automatically determine efficient representations from large unstructured
datasets, such as those arising in biology. However, empirical evidence
suggests that seemingly minor changes to these machine learning models yield
drastically different data representations that result in different biological
interpretations of data. This begs the question of what even constitutes the
most meaningful representation. Here, we approach this question for
representations of protein sequences, which have received considerable
attention in the recent literature. We explore two key contexts in which
representations naturally arise: transfer learning and interpretable learning.
In the first context, we demonstrate that several contemporary practices yield
suboptimal performance, and in the latter we demonstrate that taking
representation geometry into account significantly improves interpretability
and lets the models reveal biological information that is otherwise obscured.

## Introduction

_Data representations_ play a crucial role in the statistical analysis of
biological data, and results have lingered on appropriate choice of
representation. It is therefore not surprising to see an uprise in biology of
_representation learning_ [1], a subfield of machine learning where the
representation is estimated alongside the statistical model. In the analysis
of protein sequences in particular, the last year has produced a number of
studies that demonstrate how representations can help extract important
biological information automatically from the millions of observations
acquired through modern sequencing technologies [2, 3, 4, 5]. While these
promising results indicate that learned representations can have substantial
impact on scientific data analysis, they also beg the question: _what is a
good representation?_ It is this elementary question that we address here.

At its core, a _representation_ 111We note that ‘representation’, ‘feature’
and ‘embedding’ all seem to be used interchangeably in the literature. is a
distillation of data into an abstract and often lower dimensional space that
captures the essential features of the original data. This may subsequently be
used for data exploration, e.g. visualization, or task-specific predictions
where limited data is available. The classical principal component analysis
[6] learns features that are linearly related to the original data, while
contemporary techniques seek highly non-linear relations [1]. This approach
has been particularly successful in natural language processing (NLP), where
representations of word sequences can be learned from the vast online textual
resources, extracting general properties of language that support specific
natural language tasks [7, 4, 8]. The success of such _word sequence models_
has inspired its use for modelling _biological sequences_ , leading to
impressive results in remote homology modelling [9], prediction of function
[10], stability prediction [3], secondary structure prediction [11], etc.

Figure 1: Representations in a transfer-learning setting: a typical scenario.
In a pretraining phase, a model is trained to embed or _encode_ input protein
sequences $\textbf{S}=(s_{1},s_{2},...,s_{L})$, to a local representation
$(r_{1},r_{2},...,r_{L})$, after which it is decoded to be as similar as
possible to the original sequence. After the pre-training stage, the learned
representation can be used as a proxy for the raw input sequence, with the
hope that information thereby can be transferred from the large set of protein
sequences to a specific task of interest. When training on global properties
of proteins, the local representations $r_{i}$, are aggregated (e.g. averaged)
into a global representation which is then used as the input during training
of the task-specific model. During training on the specific task, it is
possible to also update the parameters of the encoder, thereby fine-tuning the
representation to the specific task.

Since representations are becoming an important part of biological sequence
analysis, we should think critically about whether the constructed
representations efficiently capture the information we desire. This paper aims
to start this discussion. We will focus on protein sequences, although many of
the insights should apply to other biological sequences as well. Our work
consists of two parts. First we investigate the impact of neural network
design on the resulting representation, and find that several current
practices are suboptimal. Second we present a case study demonstrating that
explicit modeling of the _representation geometry_ reveals biological
information that was otherwise hidden to the investigator. Our results
demonstrate a clear potential for _designing_ representations actively, and
for _analyzing_ them appropriately.

## Results

Representation learning has at least two uses: In _transfer learning_ we seek
a representation that improves a downstream task, and in _data interpretation_
the representation should reveal the patterns underlying data, e.g. through
visualization. Since the first has been at the center of recent literature [2,
12, 13, 3], we place our initial focus there, and return later to data
interpretation.

### Representations for transfer learning

Transfer learning address the problems caused by limited access to labelled
data. For instance, when predicting the stability of a given protein, we only
have limited training data available as it is experimentally costly to measure
stability. The key idea is to leverage the millions of available unlabelled
protein sequences to learn a general protein representation through an
_embedding model_ , and then train a problem-specific _task model_ on top
using the limited labelled training data (Figure 1).

When considering representations in the transfer-learning setting, the quality
of a representation will be judged by the level of predictive performance
obtained by one or more downstream predictive tasks. A recent study
established a benchmark set of such problems for protein sequence
representations [3]. For our experiments below, we will consider three of
these tasks, each reflecting a particular global protein property: 1)
classification of protein sequences into a set of 1,195 known folds [14], 2)
fluorescence prediction for variants of the green fluorescent protein in
_Aequorea victoria_ [15], and 3) prediction of the stability of protein
variants obtained in high throughput experimental design experiments [16].

#### Constructing a global representation as an average of local
representations is suboptimal.

For sequential data, such as proteins and DNA, the embedding model must be
constructed to process input of varying length. Sequential models have been
used for this purpose in natural language processing [17, 4, 8], and most of
the recent representation learning advances in proteins follow the same
approach [2, 12, 13, 3]. The embedding model is a large (millions of
parameters) model that aims to reproduce its own input, either by predicting
the next character given the sequence observed so far, or by predicting the
entire sequence from a partially obscured input sequence. The output of an
embedding model is then a sequence of _local representations_
$(r_{1},r_{2},...,r_{L})$ each corresponding to one amino acid in the input
sequence $\textbf{S}=(s_{1},s_{2},...,s_{L})$. In order to successfully
predict the next amino acid, $r_{i}$ should contain information about the
local neighborhood around $s_{i}$, together with some global signal from S.
The embedding model is typically trained using a large corpus of proteins
spanning multiple families.

In order to obtain a global representation of the entire protein, the variable
number of local representations must be aggregated into a fixed-size global
representation. We expect this choice to be quite critical to the nature of
the resulting representation. Standard approaches for this operation include
averaging with uniform [2] or learned attention [3] weights, but the complex
non-local interactions known to occur in a protein suggest that it could be
beneficial to allow for more complex aggregation functions. To investigate
this issue, we consider two alternative strategies:

The first strategy (Concat) avoids aggregation altogether by concatenating the
local representations $r=[r_{1},r_{2},...,r_{L},p,p,p]$ with additional
padding $p$ to adjust for sequence-length, such that $dim(r)=D$ for all $r$.
This approach preserves all information stored in $r_{i}$. To make a fair
comparison to the averaging strategy, we wish to maintain the same overall
representation size, which means that the dimensions of our local
representations $r_{i}$ must be scaled down by the maximum sequence length. In
our experiment, we implement both strategies using a ResNet-based model [18]
(see Methods), with a global representation size of 2048, leading to a local
representation size of 2048 for the average strategy, and a much smaller size
of 4 for the concatenation strategy.

Figure 2: Different strategies for obtaining global, sequence-length
independent representations (S=sequence, E=Encoder, D=Decoder and
P=Probabilities). The first two post-process the per-amino acid local
representations, while the last learns the global representation during
pretraining.

As a second strategy (Bottleneck), we investigate the possibility of
_learning_ the optimal aggregation operation. This is done with an autoencoder
[19], a simple neural network that as output predicts its own input, but
forces it through a low-dimensional bottleneck, which can be interpreted as a
global representation (Figure 2). The model thus _learns_ a generic global
representation during pre-training, in contrast to the strategies above in
which the global representation arises as a deterministic operation of learned
local representations during this initial phase.

When comparing these two aggregation strategies on the three protein
prediction tasks (Stability, Fluorescence, Remote Homology), we observe a
quite dramatic impact on performance (Table 1). The Bottleneck strategy, where
the global representation is learned, clearly outperforms the other
strategies. This was expected, since already during pre-training this model is
encouraged to find a more global structure in the representations. More
surprising are the results for the Concat strategy, as these demonstrate that
even if we restrict the local representation to be much smaller than in
standard sequential models, the fact that there is no loss of information
during aggregation has a dramatic influence on the downstream performance.
While these results might differ depending on the nature of the downstream
task, they strongly point to the potential of learning global representations
directly as part of the unsupervised pre-training procedure.

| Stability | Fluorescence | Homology
---|---|---|---
Mean | 0.42 | 0.19 | 0.27
Maximum | 0.02 | 0.02 | 0.28
Attention | 0.65 | 0.23 | 0.27
Concat | 0.74 | 0.69 | 0.34
Bottleneck | 0.79 | 0.78 | 0.41
Table 1: Comparison of strategies for obtaining global, sequence-length independent representations on three downstream tasks [3]. The first three are variants of averaging used in the literature, using uniform weights (Mean), learned attention weights (Attention), or the local representation with the highest attention weight (Maximum). The last two demonstrate that simple alternatives such as concatenating smaller local representations (Concat) or changing the model to directly learn a global representation (Bottleneck) can have a substantial impact on performance (best results in bold). | Remote Homology | Fluorescence | Stability
---|---|---|---
| Resnet | LSTM | Trans | Resnet | LSTM | Trans | Resnet | LSTM | Trans
Pre +Fix | 0.27 | 0.37 | 0.27 | 0.23 | 0.74 | 0.48 | 0.65 | 0.70 | 0.62
Pre +Fin | 0.17 | 0.26 | 0.21 | 0.21 | 0.67 | 0.68 | 0.73 | 0.69 | 0.73
Rng+Fix | 0.03 | 0.10 | 0.04 | 0.25 | 0.63 | 0.14 | 0.21 | 0.61 | -
Rng+Fin | 0.10 | 0.12 | 0.09 | -0.28 | 0.21 | 0.22 | 0.61 | 0.28 | -0.06
Baseline | 0.09 (Accuracy) | 0.14 (Correlation) | 0.19 (Correlation)
Table 2: The impact of fine-tuning and initialization on downstream model
performance. The embedding models were either randomly initialized (Rng) or
pre-trained (Pre), and subsequently either fixed (Fix) or fine-tuned to the
task (Fin). Although fine-tuning is beneficial on some task/model
combinations, we see clear signs of overfitting in the majority of cases (best
results in bold).

#### Representation learning can also overfit.

Common wisdom dictates that models overfit if they have more trainable
parameters than we have observations. While deep learning models are subject
to implicit regularization [20], we hypothesize that overfitting remains
problematic in representation learning, and here consider two specific
instances.

Conceptually, we often view the embedding and task models as being separate
(Figure 1), but it is common practice to fine-tune the embedding model for a
given task, which implies that all joint model parameters are to be optimized
[3]. To test whether this leads to overfitting, we train three models, an LSTM
[21], a Transformer [17], and a dilated residual network (ResNet) [22], where
we either keep the embedding model fixed (Fix) or fine-tune it to the task
(Fin). To evaluate the impact of the representation model itself, we consider
both a pre-trained version (Pre) and randomly initialized representation
models that are not trained on data (Rng). Such models will map similar inputs
to similar representations, but should otherwise not perform well. Finally, as
a naive baseline representation, we consider the frequency of each amino acid
in the sequence.

Table 2 shows that fine-tuning the embedding reduces test performance in two
out of three tasks, probably due to overfitting. In particular, for the best
performing model (LSTM), we observe a significant difference in performance
with and without fine-tuning. Interestingly, randomly initialized
representations perform remarkably well in several cases, which echos results
known from random projections [23]. In many cases the random models outperform
the simple one-hot encoded baseline representation. One explanation is that
even random networks map similar proteins to similar representations which may
be an important property in representation learning.

While these trends depend on the nature of the downstream task, and the amount
of available data, our experiments indicate that issues of overfitting are as
prevalent in representation learning as in all other branches of machine
learning. Fine-tuning should only then take place under rigorous cross-
validation. The results suggest that fixing the embedding model during task
training should be the default choice. In practice, there is evidence that one
can successfully fine-tune a small subset of the model parameters [24], which
may be a workable middle ground.

#### Reconstruction error is not a good measure of representation quality.

One of the driving factors of successful NLP is the increasingly larger
embedding models trained on ever-growing datasets [4, 25]. This has spread to
protein learning, in particular the prevalent idea that increasing
representation dimensionality gives better representations. For instance, the
well known UniRep model employs the largest representation allowed by the
hardware used [2]. Another potential reason for large representations is that
the choice of representation size is sometimes inconveniently tied to the
design of the embedding process itself. For instance, with an LSTM [21] the
representation is often chosen as the hidden state in one of the layers. As it
is difficult to train an LSTM with a very small hidden state, this effectively
constrains the choice of representation size, for instance making it difficult
to employ the Concat strategy introduced earlier.

Figure 3: Reconstruction and downstream performance as a function of
representation size. Although reconstruction accuracy consistently improves
for increasing representation size, the performances on the individual tasks
deteriorate for large representations. Performance refers to Spearman
correlation (stability, fluorescence) or accuracy (homology, reconstruction).

It is unclear whether the high dimensionality is beneficial to the
representation itself, or whether it merely serves to provide a lower
reconstruction error. We here explore this matter using the ResNet-based model
combined with the BottleNeck strategy (Figure 3). For the three downstream
tasks, we see that performance increases with representation size until a
tipping point, after which predictive performance drops. In contrast, the
hold-out reconstruction accuracy increases consistently with latent size.
Although other models might show different behaviors, the mismatch between
reconstruction and downstream performance observed for our example suggests
that we cannot generally rely on reconstruction error as a proxy for
representation quality.

Figure 4: Latent embedding of the protein family of $\beta$-lactamase, color-
coded by taxonomy at the phyla level. In (a), we embed the family using a
sequential model trained on the full corpus of protein families. In (b), we
have reduced the scope and train the same model on only the $\beta$-lactamase
family. In (c), we again train on just the $\beta$-lactamase family, but now
after preprocessing the sequences in a multiple sequence alignment and
applying a dense model rather than a sequentional model. Finally, in (d) we
employ a sequence reweighting scheme to compensate for the sampling bias in
sequence space. With each of these steps, we gain better separation between
the different phyla, demonstrating the impact of model choice and data
preprocessing on the learned representation.

### Representations for data interpretation

If a representation accurately describes structure in the underlying dataset,
we might expect it to be useful not only as input to a downstream model, but
also as the basis for direct interpretation, for instance through
visualization. In this context, it is important to realize that different
modelling choices can lead to dramatically different representations of the
same data. More troubling, even when using the same model assumptions,
repeated training instances can also deviate substantially, and we must
therefore analyze our interpretations with care. In the following, we explore
these effects in detail.

#### Modeling choices have a dramatic effect on learned representations.

Recent generative models for proteins tend to learn universal representations
of protein space akin to language models. In bioinformatics, there is,
however, a long history of analyzing protein per family. Since the proteins in
the same family share a common three-dimensional structure, an underlying
correspondence exists between positions in different sequences, which we can
approximate using multiple sequence alignment techniques. After establishing
such an alignment, all input sequences will have the same length, making it
possible to use simple fixed-size input models, rather than the sequential
models discussed previously. One advantage is that models can now readily
detect patterns at and correlations between absolute positions of the input,
and directly observe both conservation and coevolution. In terms of
interpretability, this has clear advantages. An example of this approach is
the DeepSequence model [26], in which the latent space of a VAE was shown to
clearly separate the input sequences into different phyla, and capture
covariance among sites on par with earlier coevolution methods. We reproduce
this result here using a VAE on the $\beta$-lactamase family (Figure 4c).

If we apply the globally trained sequence models from the previous sections on
the same set of proteins from the $\beta$-lactamase family, we see no such
phylogenetic separation (Figure 4a). The fact that the phyla are much less
clearly resolved in the sequence model is perhaps unsurprising, since this
representation has been trained to represent the space of all proteins, and
therefore does not have the same capacity to separate details of a single
protein family. Indeed, to compensate for this, recent work has introduced the
concept of evo-tuning, where a global representation is fine-tuned on a single
protein family [2]. However, when training exclusively on $\beta$-lactamase
(Figure 4b), we still do not see the same structure as in the alignment-based
VAE model (Figure 4c). This is despite the fact that the two models now are
trained on the same set of proteins.

The results above can be explained either by the inherent differences in the
model architectures, or by the domain-specific knowledge inserted through
preprocessing sequences when constructing an alignment. The different
inductive biases in the two model architectures certainly play a role. For
instance, sequential models generally still struggle to recover covariance
signals in unaligned protein sequences, although recent progress has been made
in this area [27]. However, one additional factor that is sometimes overlooked
is that aligned sequences also facilitate further preprocessing of the input
data, leading to potential benefits in performance. For instance, as is common
in coevolution analysis [28], the DeepSequence model weights input sequences
by the number of neighboring sequences within a fixed distance, thus
compensating for a potential sampling bias in sequence space. This
preprocessing was shown to have a substantial effect on change-of-stability
prediction performance [26]. Implementing this reweighting technique in our
example, we also see visually that it has a significant impact on the obtained
representation (Figure 4d). This suggests that density reweighting may be of
more general value in generative models of proteins.

Figure 5: A phylogenetic tree encoded into the latent representation space.
The representation and colors correspond to Fig. 4c. The internal nodes were
determined using ancestral reconstruction after inferring a phylogenetic tree
(branches encoded in black, leaf-nodes in gray). Figure 6: Correlation between
distances in latent space and phylogenetic distances.

#### Geometry gives meaning to representations.

Perhaps the most exciting prospect of representation learning is the
possibility of gaining new insights through interpretation and manipulation of
the learned representation space. In NLP, the celebrated word2vec model [29]
demonstrated that simple arithmetic operations on representations yielded
meaningful results, e.g. ”Paris - France + Italy = Rome”. The ability to
perform such operations on proteins would have substantial impact on protein
engineering and design. However, the lack of such results in the literature
suggests that some aspects of learned representations remain poorly
understood.

To qualify the discussion, we note that standard arithmetic operations such as
addition and subtraction relies on the assumption that the learned
representation space is Euclidean. The visualizations of the alignment-based
family-specific VAE representation in Fig. 4, however, suggest that a
Euclidean interpretation is misleading: The star-like structure of the
corresponding representations resembles the structure of a phylogenetic tree
rooted at the origin. To investigate this structural origin of the latent
space, we estimate a phylogenetic tree of a subset of our input data, and
encode the inner nodes of the tree to our latent space using ancestral
reconstruction (see Methods). This confirms that the learned representation to
a large extent follows the evolutionary history behind the data (Fig. 5). If
we define similarities between pairs of points through the Euclidean distance
between them, we would obtain straight-line interpolants that pass through
uncharted territory in the representation space. This does not seem fruitful.

Empirically, we also observe that the Euclidean interpretation does not
reflect the driving evolution. The first two panels of Fig. 6 show the
correlation between Euclidean distances and phylogenetic distances in an LSTM
and a VAE representation. We observe practically no correlation in the LSTM
representation, while the VAE representation shows some correlation. The
slight positive correlation obtained in the VAE may be explained by Euclidean
distances to the origin being correlated with phylogenetic distance [30].

Mathematically, the Euclidean interpretation is problematic as well. In
general, the latent variables of a generative model are not statistically
identifiable, such that it is most often possible to deform the latent
representation space without changing the estimated data density [31, 32].
With this in mind, the Euclidean assumption seems difficult to justify beyond
arguments of simplicity, as Euclidean arithmetic is not invariant to general
deformations of the representation space. It has recently been pointed out
that shortest paths (geodesics) and distances between representation pairs can
be made identifiable even if the latent coordinates of the points themselves
are not [33, 32]. The trick is to equip the learned representation with a
Riemannian metric which ensures that distances are measured in data space
along the estimated manifold. This result suggests that perhaps a Riemannian
set of operations is more suitable for interacting with learned
representations than the usual Euclidean arithmetic operators.

Figure 7: Shortest paths (geodesics) between representations of
$\beta$-lactamase in a VAE. The Riemannian metric corresponds to measuring the
expected distance between one-hot encoded proteins measured along the
estimated manifold. The geodesics generally move along the star-shaped
structure of the data similarly to the estimated phylogenetic tree, suggesting
that the geodesics are well-suited for interpolating proteins.

To investigate this hypothesis, we first develop a suitable Riemannian metric,
such that geodesic distances correspond to expected distances between one-hot
encoded proteins, which are integrated along the manifold. The details of this
construction are provided in the Methods section. The last panel of Fig. 6
shows the correlation between such geodesic distances and phylogenetic
distances. Unlike in the Euclidean interpretation, we now observe a
significant correlation. For short-to-medium distances the correlation is near
perfect. This suggests that the Riemannian interpretation of the
representation is more in tune with the driving biology.

Visually, the difference is also quite striking (Fig. 7). We see that the
geodesics follow the manifold very closely, and to a large extent preserve the
underlying tree structure. As was already apparent in Fig. 4c, the
organization of clusters in latent space is irregular. For instance the yellow
subtree in the top right corner would seem to fit better in the bottom left,
suggesting a suboptimal fit of the original VAE. We see that both the
phylogenetic reconstruction and our geodesics agree on this, which is visually
clear by the thick bundle of geodesics running diagonally to connect these
regions. The structures closest to the leaves (i.e. the local differences)
seem to be subject to most discrepancies between the phylogenetic
reconstruction and the geodesics. Which of the two gives the most accurate
description is not clear. Since we use a fairly rough phylogenetic
reconstruction, and the VAE encodes more complex correlation between sites,
one could speculate that the VAE latent space might provide a richer
description of similarity, but more research is required to settle this
matter.

## Discussion

Learned representations of protein sequences have the potential to
substantially improve systems for making biological predictions, and may also
help to reveal previously uncovered biological information. In this paper, we
have illuminated parts of the answer to the titular question of what
constitutes a good representation of proteins, but do not provide the full
answer. We first note that the question itself is vague and should be
countered with _for what should the representation be used?_ A representation
that is suitable for making predictions may not be optimal for a human
investigator to better understand the underlying biology, and vice versa. The
enticing idea of a single protein representation for all tasks might therefore
prove unworkable in practice.

### Designing purposeful representations

To design a representation for a given task requires reflection over which
biological properties we wish the representation to encapsulate. Different
biological aspects of a protein will place different demands on the
representations, but it is not straightforward to enforce specific properties
in a representation. We may, however, steer the representation learning in the
right direction by 1) picking appropriate model architectures, 2)
preprocessing the data, 3) choosing suitable objective functions, and 4)
placing prior distributions on parts of the model. We discuss each of these in
turn.

Informed network architectures can be difficult to construct as the usual
neural network ‘building blocks’ are fairly elementary mathematical functions
that are not immediately linked to high-level biological information.
Nonetheless, our discussion of length-invariant sequence representations is a
simple example of how one might inform the model architecture of the biology
of the task. It is generally acknowledged that global protein properties are
not linearly related to local properties. It is therefore not surprising when
we show that the model performance significantly improves when we allow the
model to learn such a nonlinear relationship instead of relying on the common
linear average of local representations. We speculate that while similar ‘low-
hanging fruit’ may remain in currently applied network architectures, they are
limited, and more advanced tools are needed to encode biological information
into network architectures. The internal representations in attention-based
architectures have been shown to recover known physical interactions between
proteins [34], opening the door to the incorporation of prior information
about known physical interactions in a protein. Recent work on permutation and
rotation invariance/equivariance in neural networks [35, 36] may also prove
valuable for this purpose, though they have yet to be explored exhaustively in
representation learning.

Data preprocessing is frowned upon in contemporary ‘end-to-end‘ representation
learning, but it remains an important part of model design. In particular,
preprocessing using the vast selection of existing tools from computational
biology is a valuable way to encode existing biological knowledge into the
representation. We saw a significant improvement in the representation
capabilities of unsupervised models when trained on aligned protein sequences,
as this injects prior knowledge about comparable sequence positions in a set
of sequences. While recent work is increasingly working towards techniques for
learning such signals directly from data [37], these approaches have yet to
reach a mature state. This example suggests that if we move too fast towards
‘end-to-end’ learning, we may run the risk of throwing the baby out with the
bathwater, by discarding years of experience endowed in existing tools.

Relevant objective functions are paramount to any learning task. Although
representation learning is typically conducted using a reconstruction loss, we
demonstrate that optimal representations according to this objective are
generally sub-optimal for any specific task at hand (Fig. 3). This suggests
that it is difficult to optimize our representation models without knowing the
downstream task. In addition, multiple different reconstruction objectives
have recently come out of the natural language processing field (e.g. next-
character prediction and various masked input strategies), and it remains
unclear how the resulting representations depend on these choices. This
suggests that hyper-parameters of representations should be chosen based on
downstream task-specific performance, rather than reconstruction performance
on a hold-out set. This is, however, a delicate process, as optimizing the
actual _parameters_ of the representation model on the downstream task is
associated with a high risk of overfitting. We anticipate that principled
techniques for combining reconstruction objectives on the large unsupervised
data sets with task specific objectives in a semi-supervised learning setting
could provide substantial benefits in this area [38, 9].

Informative priors can be useful to impose softer preferences than those
encoded by hard architecture constraints. The standard Gaussian prior in the
variational autoencoder is such an example, though its preference is not
guided by biological information, which appears to be a missed opportunity. In
the studies of $\beta$-lactamase, we, and others, observe a representation
structure that somewhat resembles the phylogenetic tree spanned by the
evolution of the protein family. We speculate that recent hyperbolic priors
[39] that are designed to emphasize hierarchies in data may bring forward such
evolutionary structure. Since we observe that the latent representation better
reflects biology when endowed with a suitable Riemannian metric, it may also
be valuable to leverage this geometry directly in the prior [40].

### Analysing representations appropriately

Even with the most valiant efforts to incorporate prior knowledge into our
representations, they must still be interpreted with great care. We highlight
the particular example of distances in representation space, and emphasize
that the seemingly natural Euclidean distances are misleading. The non-
linearity of encoders and decoders in modern machine learning methods means
that representation spaces are generally non-Euclidean. We have demonstrated
that by bringing the expected distance from the observation space into the
representation space in the form of a Riemannian metric, we obtain geodesic
distances that correlate significantly better with phylogenetic distances than
what can be attained through the usual Euclidean view. This is an exciting
result as the Riemannian view comes with a set of natural operators akin to
addition and subtraction, such that the representation can be engaged with
operationally. We expect this to be valuable for e.g. protein engineering,
since it gives an operational way to combine representations from different
proteins.

The geometric analysis comes with several implications that are relevant
beyond proteins. It suggests that the commonly applied visualizations where
latent representations are plotted as points on a Euclidean screen may be
highly misleading. We therefore see a need for visualization techniques that
faithfully reflect the geometry of the representations. The analysis also
indicates that downstream prediction tasks may gain from leveraging the
geometry, although standard neural network architectures do not yet have such
capabilities.

## Methods

#### Variational autoencoders.

A variational autoencoder assumes that data X is generated from some (unknown)
latent factors Z though the process $p_{\theta}(\textbf{X}|\textbf{Z})$. The
latent variables Z can be viewed as the compressed representation of X. Latent
space models try to model the joint distribution of X and Z as
$p_{\theta}(\textbf{X},\textbf{Z})=p_{\theta}(\textbf{Z})p_{\theta}(\textbf{X}|\textbf{Z})$.
The generating process can then be viewed as a two step procedure: first a
latent variable Z is sampled from the prior and then data X is sampled from
the conditional $p_{\theta}(\textbf{X}|\textbf{Z})$ (often called the
decoder). Since X is discrete by nature, $p_{\theta}(\textbf{X}|\textbf{Z})$
is modelled as a Categorical distribution
$p_{\theta}(\textbf{X}|\textbf{Z})\sim Cat(C,l_{\theta}(\textbf{Z}))$ with $C$
classes and $l_{\theta}(Z)$ being the log-probabilities for each class. To
make the model flexible enough to capture higher order amino acid
interactions, we model $l_{\theta}(Z)$ as a neural network. Even though data X
is discrete, we use continuous latent variables $\textbf{Z}\sim N(0,1)$.

#### Distance in sequence space.

To calculate geodesic distances we first need to define geodesics over the
random manifold defined by $p(\textbf{X}|\textbf{Z})$. These geodesics are
curves $c$ that minimize expected _energy_ [32] defined as

$E(c)=\int
E\left[\left\lVert\partial_{t}f(c_{t})\right\rVert^{2}\right]\>\mathrm{d}t,$
(1)

where $f$ is the stochastic mapping from latent space to data space. This
energy requires a meaningful (squared) norm in data space. We remind here that
protein sequence data $x,y$ is embedded into a one-hot space i.e.

$x,y=\left\\{\begin{bmatrix}1\\\ 0\\\ \vdots\\\
0\end{bmatrix},\begin{bmatrix}0\\\ 1\\\ \vdots\\\
0\end{bmatrix},\begin{bmatrix}0\\\ 0\\\ \vdots\\\ 1\end{bmatrix}\right\\},$

where we assume that $p(x_{d}=1)=a_{d}$, $p(y_{d}=1)=b_{d}$ for $d=1,...,C$.
It can easily be shown that the squared norm between two such one-hot vectors
can either be 0 or 2:

$\Delta^{2}=\left\lVert x-y\right\rVert^{2}=\\{0,2\\}.$

The probability of these two events are given as

$\displaystyle P(\Delta^{2}=0)$ $\displaystyle=P(x=y)$
$\displaystyle=P(x_{1}=y_{1})+P(x_{2}=y_{2})+...+P(x_{D}=y_{D})$
$\displaystyle=\sum_{d=1}^{D}a_{d}b_{d}$ $\displaystyle P(\Delta^{2}=2)$
$\displaystyle=1-P(\Delta^{2}=0)=1-\sum_{d=1}^{D}a_{d}b_{d}.$

The expected squared distance is then given by

$\displaystyle\mathop{\mathbb{E}}(\Delta^{2})$
$\displaystyle=\int_{\\{0,2\\}}\Delta^{2}\cdot P(\Delta^{2})d\Delta^{2}$
$\displaystyle=0\cdot P(\Delta^{2}=0)+2\cdot P(\Delta^{2}=2)$
$\displaystyle=2\left(1-\sum_{d=1}^{D}a_{d}b_{d}\right),$

Extending this measure to two sequences of length $L$ is then

$\mathop{\mathbb{E}}(\Delta^{2})=\sum_{l=1}^{L}2\left(1-\sum_{d=1}^{D}a_{l,d}b_{l,d}\right).$
(2)

The curve energy (1) can then be approximated as

$E(c)\approx\sum_{i=1}^{N-1}\sum_{l=1}^{L}2\left(1-\sum_{d=1}^{D}p(c_{i})_{l,d}p(c_{i+1})_{l,d}\right)\left\lVert
c_{i+1}-c_{i}\right\rVert_{2}^{2}.$

We note that the discrete energy is a trade-off between having low distance in
data space and low distance in latent space.

#### Construction of entropy network.

To build a reliable Riemannian metric, we need the expected curve energy favor
regions of high data density. It has been shown that this require a well-
calibrated measure of uncertainty [32, 41]. That is, the network
$p_{\theta}(\textbf{X}|\textbf{Z})$ should be certain about its output in
regions where we have observed data, and uncertain in regions where we have
not. In a standard VAE with posterior modelled as a normal distribution
$\mathcal{N}(\mu_{\theta}(\textbf{Z}),\sigma^{2}_{\theta}(\textbf{Z}))$, this
amounts to constructing a variance network $\sigma^{2}_{\theta}(\textbf{Z})$
that increases away from data [33, 42]. However, no prior work has been done
on discrete distributions, such as the Categorical distribution
$C(\mu_{\theta}(\textbf{Z}))$ that we are working with. In this model we do
not have a clear division of the average output (mean) and uncertainty
(variance), so we control the uncertainty through the entropy of the
distribution. We remind that for a categorical distribution, the entropy is

$H(\textbf{X}|\textbf{Z})=\sum_{i=1}^{C}p_{\theta}(\textbf{X}|\textbf{Z})_{i}\cdot\log
p_{\theta}(\textbf{X}|\textbf{Z})_{i}.$

The most uncertain case corresponds to when $H(\textbf{X}|\textbf{Z})$ is
largest i.e. when $p(\textbf{X}|\textbf{Z})_{i}=1/C$ for $i=1,...,C$. Thus, we
want to construct a network $p_{\theta}(\textbf{X}|\textbf{Z})$ that assigns
equal probability to all classes when we are away from data, but is still
flexible when we are close to data. Taking inspiration from [42] we construct
a function $\alpha=T(z)$, that maps distance in latent space to the zero-one
domain ($T:[0,\inf)\mapsto[0,1]$). $T$ is a trainable network of the model,
with the functional form $T(z)=\texttt{sigmoid}(\beta\cdot V(z))$ with
$V(z)=\min_{j=\\{1,..,K\\}}\left\lVert z-c_{j}\right\rVert_{2}^{2}$, where
$c_{j}$ are trainable cluster centers (initialized using $k$-means). This
function essentially estimates how close a latent point $z$ is to the data
manifold, returning 1 if we are close and 0 when far away. Here $K$ indicates
the number of cluster centers (hyperparameter) and $\beta$ is a overall
scaling (trainable). With this network we can ensure a well-calibrated entropy
by picking

$p_{\theta}(\textbf{X}|\textbf{Z})_{i}=\alpha\cdot
p_{\theta}(\textbf{X}|\textbf{Z})_{i}+(1-\alpha)\cdot\mathbb{L},$

where $\mathbb{L}=\frac{\sum_{i=1}^{C}p_{\theta}(\textbf{X}|\textbf{Z})}{C}$.
For points far away from data, we have $\alpha=0$ and return $\mathbb{L}$
regardless of category (class), giving maximal entropy. When near the data, we
have $\alpha=1$ and the entropy is determined by the trained decoder
$p_{\theta}(\textbf{X}|\textbf{Z})_{i}$.

Figure 8 shows the difference in entropy of the likelihood between a standard
VAE (top) and a VAE equipped with our developed entropy network (bottom). The
standard VAEs produce arbitrary entropy, and is often more confident in its
predictions far away from the data. Our network increases entropy as we move
away from data.

Figure 8: Construction of the entropy network for our geodesic calculations.
Top: latent representations of $\beta$-lactamase with the background color
denoting the entropy of the output posterior for a standard VAE. Bottom: as
top but using a VAE equipped with our developed entropy network.

#### Optimizing geodesics.

In principal, the geodesics could be found by direct minimization of the
expected energy. However, empirically we observed that this strategy was prone
to diverge, since the optimization landscape is very flat near the initial
starting point. We therefore instead discretise the entropy landscape into a
2D grid, and form a graph based on this. In this graph each node will be a
point in the grid, which is connected to its eight nearest neighbours, with
the edge weight being the distance. Then using Dijkstra’s algorithm [43] we
can very fast get a very robust initialization of each geodesic. To get the
final geodesic curve we fit a cubic spline [44] to the discretized curve found
by Dijkstra’s algorithm, and afterwards do 10 gradient steps to refine the
solution.

#### Phylogeny and ancestral reconstruction

We used FastTree2 [45] with standard settings for estimation of phylogenetic
trees and subsequently applied the codeml program [46] from the PAML package
for ancestral reconstruction of the internal nodes of the tree.

#### Acknowledgements.

This project has received funding from the European Research Council (ERC)
under the European Union’s Horizon 2020 research and innovation programme
(grant agreement no 757360). NSD and SH were supported in part by a research
grant (15334) from VILLUM FONDEN. W.B. was supported by a project grant from
the Novo Nordisk Foundation (NNF18OC0052719). We gratefully acknowledge the
support of NVIDIA Corporation with the donation of GPU hardware used for this
research.

#### Author contributions.

N.S.D., S.H. and W.B. jointly conceived and designed the study. N.S.D
conducted the experiments. All authors contributed to the writing of the
paper.

## References

  * [1] Y. Bengio, A. Courville and P. Vincent “Representation Learning: A Review and New Perspectives” In _IEEE Transactions on Pattern Analysis and Machine Intelligence_ 35.8, 2013, pp. 1798–1828 DOI: 10.1109/TPAMI.2013.50
  * [2] Ethan C Alley, Grigory Khimulya, Surojit Biswas, Mohammed AlQuraishi and George M Church “Unified rational protein engineering with sequence-based deep representation learning” In _Nature methods_ 16.12 Nature Publishing Group, 2019, pp. 1315–1322
  * [3] Roshan Rao et al. “Evaluating protein transfer learning with TAPE” In _Advances in Neural Information Processing Systems_ , 2019, pp. 9689–9701
  * [4] Jacob Devlin, Ming-Wei Chang, Kenton Lee and Kristina Toutanova “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding” In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , 2019, pp. 4171–4186
  * [5] Alexander Rives et al. “Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences” In _bioRxiv_ Cold Spring Harbor Laboratory, 2019, pp. 622803
  * [6] I.T. Jolliffe “Principal Component Analysis” Springer, 1986
  * [7] Alec Radford, Karthik Narasimhan, Tim Salimans and Ilya Sutskever “Improving Language Understanding by Generative Pre-Training”, 2018
  * [8] Yinhan Liu et al. “RoBERTa: A Robustly Optimized BERT Pretraining Approach”, 2019 arXiv:1907.11692
  * [9] Seonwoo Min, Seunghyun Park, Siwon Kim, Hyun-Soo Choi and Sungroh Yoon “Pre-Training of Deep Bidirectional Protein Sequence Representations with Structural Information”, 2019 arXiv:1912.05625
  * [10] Maxat Kulmanov, Robert Hoehndorf and Lenore Cowen “DeepGOPlus: Improved protein function prediction from sequence” In _Bioinformatics_ 36.2 Oxford University Press, 2020, pp. 422–429 DOI: 10.1093/bioinformatics/btz595
  * [11] Mirko Torrisi, Manaz Kaleel and Gianluca Pollastri “Deeper Profiles and Cascaded Recurrent and Convolutional Neural Networks for state-of-the-art Protein Secondary Structure Prediction” In _Sci. Rep._ 9.1 Nature Publishing Group, 2019, pp. 1–12 DOI: 10.1038/s41598-019-48786-x
  * [12] Jose Juan Almagro Armenteros, Alexander Rosenberg Johansen, Ole Winther and Henrik Nielsen “Language modelling for biological sequences–curated datasets and baselines” In _bioRxiv_ Cold Spring Harbor Laboratory, 2020
  * [13] Michael Heinzinger et al. “Modeling aspects of the language of life through transfer-learning protein sequences” In _BMC bioinformatics_ 20.1 Springer, 2019, pp. 723
  * [14] Jie Hou, Badri Adhikari and Jianlin Cheng “DeepSF: deep convolutional neural network for mapping protein sequences to folds.” In _Bioinformatics_ 34.8 Oxford University Press, 2018, pp. 1295–1303 DOI: 10.1093/bioinformatics/btx780
  * [15] Karen S. Sarkisyan et al. “Local fitness landscape of the green fluorescent protein” In _Nature_ 533.7603 Nature Publishing Group, 2016, pp. 397–401
  * [16] Gabriel J. Rocklin et al. “Global analysis of protein folding using massively parallel design, synthesis, and testing” In _Science_ 357.6347 American Association for the Advancement of Science, 2017, pp. 168–175
  * [17] Ashish Vaswani et al. “Attention is all you need” In _Advances in neural information processing systems_ , 2017, pp. 5998–6008
  * [18] Kaiming He, Xiangyu Zhang, Shaoqing Ren and Jian Sun “Deep residual learning for image recognition” In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 770–778
  * [19] Mark A. Kramer “Nonlinear principal component analysis using autoassociative neural networks” In _AIChE Journal_ 37.2 John Wiley & Sons, Ltd, 1991, pp. 233–243 DOI: 10.1002/aic.690370209
  * [20] Tomaso Poggio, Qianli Liao and Andrzej Banburski “Complexity control by gradient descent in deep networks” In _Nature communications_ 11.1 Nature Publishing Group, 2020, pp. 1–5
  * [21] Sepp Hochreiter and Jürgen Schmidhuber “Long Short-Term Memory” In _Neural Compututation_ 9.8 MIT Press Journals, 1997, pp. 1735–1780 DOI: 10.1162/neco.1997.9.8.1735
  * [22] Fisher Yu, Vladlen Koltun and Thomas Funkhouser “Dilated residual networks” In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2017, pp. 472–480
  * [23] Ella Bingham and Heikki Mannila “Random projection in dimensionality reduction: applications to image and text data” In _Proceedings of the seventh ACM SIGKDD international conference on Knowledge discovery and data mining_ , 2001, pp. 245–250
  * [24] Nils Strodthoff, Patrick Wagner, Markus Wenzel and Wojciech Samek “UDSMProt: universal deep sequence models for protein classification” In _Bioinformatics_ 36.8, 2020, pp. 2401–2409 DOI: 10.1093/bioinformatics/btaa003
  * [25] Tom B. Brown et al. “Language Models are Few-Shot Learners”, 2020 arXiv: http://arxiv.org/abs/2005.14165
  * [26] Adam J. Riesselman, John B. Ingraham and Debora S. Marks “Deep generative models of genetic variation capture the effects of mutations” In _Nature Methods_ 15.10 Nature Publishing Group, 2018, pp. 816–822 DOI: 10.1038/s41592-018-0138-4
  * [27] Adam Riesselman et al. “Accelerating Protein Design Using Autoregressive Generative Models” In _bioRxiv_ Cold Spring Harbor Laboratory, 2019, pp. 757252 DOI: 10.1101/757252
  * [28] Magnus Ekeberg, Cecilia Lövkvist, Yueheng Lan, Martin Weigt and Erik Aurell “Improved contact prediction in proteins: Using pseudolikelihoods to infer Potts models” In _Physical Review E_ 87.1 American Physical Society (APS), 2013 DOI: 10.1103/physreve.87.012707
  * [29] Tomas Mikolov, Kai Chen, Greg Corrado and Jeffrey Dean “Efficient estimation of word representations in vector space”, 2013 arXiv:1301.3781
  * [30] Xinqiang Ding, Zhengting Zou and Charles L. Brooks “Deciphering protein evolution and fitness landscapes with latent space models” In _Nature Communications_ 10.1 Nature Research, 2019 DOI: 10.1038/s41467-019-13633-0
  * [31] Christopher M. Bishop “Pattern Recognition and Machine Learning” Springer, 2006
  * [32] Søren Hauberg “Only Bayes should learn a manifold (on the estimation of differential geometric structure from data)”, 2018 arXiv: http://arxiv.org/abs/1806.04994
  * [33] Georgios Arvanitidis, Lars Kai Hansen and Søren Hauberg “Latent space oddity: On the curvature of deep generative models” In _6th International Conference on Learning Representations, ICLR 2018_ , 2018
  * [34] Jesse Vig, Ali Madani, Lav R Varshney, Caiming Xiong, Richard Socher and Nazneen Fatema Rajani “BERTology Meets Biology: Interpreting Attention in Protein Language Models”, 2020 arXiv:2006.15222
  * [35] Taco S. Cohen, Mario Geiger and Maurice Weiler “Intertwiners between Induced Representations (with Applications to the Theory of Equivariant Neural Networks)”, 2018 arXiv:1803.10743 [cs.LG]
  * [36] Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma and Taco S Cohen “3D Steerable CNNs: Learning rotationally equivariant features in volumetric data” In _Advances in Neural Information Processing Systems_ , 2018
  * [37] Tristan Bepler and Bonnie Berger “Learning protein sequence embeddings using information from structure” In _International Conference on Learning Representations_ , 2019
  * [38] Amy X. Lu, Haoran Zhang, Marzyeh Ghassemi and Alan Moses “Self-Supervised Contrastive Learning of Protein Representations By Mutual Information Maximization” Cold Spring Harbor Laboratory, 2020 DOI: 10.1101/2020.09.04.283929
  * [39] Emile Mathieu, Charline Le Lan, Chris J Maddison, Ryota Tomioka and Yee Whye Teh “Continuous hierarchical representations with poincaré variational auto-encoders” In _Advances in neural information processing systems_ , 2019
  * [40] Dimitris Kalatzis, David Eklund, Georgios Arvanitidis and Søren Hauberg “Variational Autoencoders with Riemannian Brownian Motion Priors” In _Proceedings of the 37th International Conference on Machine Learning_ 119, 2020
  * [41] Alessandra Tosi, Søren Hauberg, Alfredo Vellido and Neil D. Lawrence “Metrics for Probabilistic Geometries” In _The Conference on Uncertainty in Artificial Intelligence (UAI)_ , 2014
  * [42] Nicki Skafte, Martin Jørgensen and Søren Hauberg “Reliable training and estimation of variance networks” In _Advances in Neural Information Processing Systems_ , 2019
  * [43] Edsger W Dijkstra “A note on two problems in connexion with graphs” In _Numerische mathematik_ 1.1, 1959, pp. 269–271
  * [44] E.N. J.H. and J.. Walsh “The theory of splines and their applications” In _Canadian Mathematical Bulletin_ 11.3 Cambridge University Press, 1968, pp. 507–508 DOI: 10.1017/S0008439500029659
  * [45] Morgan N. Price, Paramvir S. Dehal and Adam P. Arkin “FastTree 2 - Approximately maximum-likelihood trees for large alignments” In _PLoS One_ 5.3 Public Library of Science, 2010, pp. e9490 DOI: 10.1371/journal.pone.0009490
  * [46] Jun Adachi and Masami Hasegawa “MOLPHY version 2.3: programs for molecular phylogenetics based on maximum likelihood” Institute of Statistical Mathematics Tokyo, 1996


11institutetext: IBM Research Almaden, San Jose, CA, USA 22institutetext: IBM
Research, Zurich, Switzerland 33institutetext: Department of Radiology, VA
Southern Nevada Healthcare System, NV, USA

# MaxCorrMGNN: A Multi-Graph Neural Network Framework for Generalized
Multimodal Fusion of Medical Data for Outcome Prediction

Niharika S. D’Souza 11 Hongzhi Wang 11 Andrea Giovannini 22 Antonio
Foncubierta-Rodriguez 22 Kristen L. Beck 11 Orest Boyko 33 Tanveer Syeda-
Mahmood 11

###### Abstract

With the emergence of multimodal electronic health records, the evidence for
an outcome may be captured across multiple modalities ranging from clinical to
imaging and genomic data. Predicting outcomes effectively requires fusion
frameworks capable of modeling fine-grained and multi-faceted complex
interactions between modality features within and across patients. We develop
an innovative fusion approach called MaxCorr MGNN that models non-linear
modality correlations within and across patients through Hirschfeld-Gebelein-
Reǹyi maximal correlation (MaxCorr) embeddings, resulting in a multi-layered
graph that preserves the identities of the modalities and patients. We then
design, for the first time, a generalized multi-layered graph neural network
(MGNN) for task-informed reasoning in multi-layered graphs, that learns the
parameters defining patient-modality graph connectivity and message passing in
an end-to-end fashion. We evaluate our model an outcome prediction task on a
Tuberculosis (TB) dataset consistently outperforming several state-of-the-art
neural, graph-based and traditional fusion techniques.

###### Keywords:

Multimodal Fusion Hirschfeld-Gebelein-Reǹyi (HGR) maximal correlation Multi-
Layered Graphs Multi-Graph Neural Networks

##  1 Introduction

In the age of modern medicine, it is now possible to capture information about
a patient through multiple data-rich modalities to give a holistic view of a
patient’s condition. In complex diseases such as cancer [18], tuberculosis
[16] or autism spectrum disorder [9, 8, 7], evidence for a diagnosis or
treatment outcome may be present in multiple modalities such as clinical,
genomic, molecular, pathological and radiological imaging. Reliable patient-
tailored outcome prediction requires fusing information from modality data
both within and across patients. This can be achieved by effectively modeling
the fine-grained and multi-faceted complex interactions between modality
features. In general, this is a challenging problem as it is largely unclear
what information is best captured by each modality, how best to combine
modalities, and how to effectively extract predictive patterns from data [14].

###  1.1 Related Works

Existing attempts to fuse modalities for outcome prediction can be divided
into at least three approaches, namely, feature vector-based, statistical or
graph-based approaches. The vector-based approaches perform early,
intermediate, or late fusion [18, 1] with the late fusion approach combining
the results of prediction rather than fusing the modality features. Due to the
restrictive nature of the underlying assumptions, these are often inadequate
for characterizing the broader range of relationships among modality features
and their relevance for prediction. In statistical approaches, methods such as
canonical correlation analysis [19] and its deep learning variants [25]
directly model feature correlations either in the native representation or in
a latent space [18]. However, these are not guaranteed to learn discriminative
patterns in the unsupervised setting and can suffer from scalability issues
when integrated into larger predictive models [23]. Recently, graph-based
approaches have been developed which form basic [2, 26, 5] or multiplexed
graphs [6] from latent embeddings derived from modality features using
concatenation [2] or weighted averaging [26]. Task-specific fusion is then
achieved through inference via message passing walks between nodes in a graph
neural network. In the basic collapsed graph construction, the inter-patient
and intra-patient modality correlations are not fully distinguished.
Conversely, in the multiplexed formulation [6], only a restricted form of
multi-relational dependence is captured between nodes through vertical
connections. Since the graph is defined using latent embedding directions, the
modality semantics are not preserved. Additionally, the staged training of the
graph construction and inference networks do not guarantee that the
constructed graphs retain discriminable interaction patterns.

###  1.2 Our Contributions

We develop a novel end-to-end fusion framework that addresses the limitations
mentioned above. The Maximal Correlation Multi-Layered Graph Neural Network,
i.e. MaxCorrMGNN, is a general yet interpretable framework for problems of
multimodal fusion with unstructured data. Specifically, our approach marries
the design principles of statistical representation learning with deep
learning models for reasoning from multi-graphs.

The main contributions of this work are three-fold:

  * •

First, we propose to model intra and inter-patient modality relationships
explicitly through a novel patient-modality multi-layered graph as shown in
Fig. 1. The edges in each layer (plane of the multi-graph) capture the intra-
modality relations between patients, while the cross-edges between layers
capture inter-modality relations across patients.

  * •

Since these relationships are not known apriori for unstructured data, we
propose, for the first time, to use learnable Hirschfeld-Gebelein-Reǹyi (HGR)
maximal correlations. We introduce learnable soft-thresholding to uncover
salient connectivity patterns automatically. Effectively, this procedure
allows us to express any multimodal dataset as a patient-modality multilayered
graph for fusion.

  * •

Third, we develop a multilayered graph neural network (MGNN) from first
principles for task-informed reasoning from multi-layered graphs.

To demonstrate the generality of our approach, we evaluate our framework on a
large Tuberculosis (TB) dataset for multi-outcome prediction. Through rigorous
experimentation, we show our framework outperforms several state-of-the-art
graph based and traditional fusion baselines.

Figure 1: Our Generalized Framework for Multimodal Fusion Gray Box: The
features from different modalities are input to the MaxCorr (HGR) formulation.
The nodes are the patients of the Multi-Graph, and the planes are the
modalities. Purple Box: The Multi-Graph Neural Network maps the multi-graph
representation to the targets.

##  2 MaxCorrMGNN formulation for Multimodal Fusion

We now describe the four main aspects of our formulation, namely, (a)
multilayered graph representation, (b) formalism for maximal correlation (c)
task-specific inference through graph neural networks, and (d) loss function
for end-to-end learning of both graph connectivity and inference.

###  2.1 Patient-Modality Multi-layered Graph:

Given multimodal data about patients, we model the modality and patient
information through a multi-layered graph [3] as shown in Fig. 1. Here the
nodes are grouped into multiple planes, each plane representing edge-
connectivity according to an individual modality while each patient is
represented by a set of corresponding nodes across the layers (called a supra-
node).

Mathematically, we represent the multi-layered graph as:
$\mathcal{G}_{\text{M}}=(\mathcal{V},\mathcal{E}_{\text{M}})$, where
$|{\mathcal{V}_{\text{M}}}|=|{\mathcal{V}}|\times K$ are the extended supra-
nodes and
$\mathcal{E}_{\text{M}}=\\{(i,j)\in\mathcal{V}_{\text{M}}\times\mathcal{V}_{\text{M}}\\}$
are the edges between supra-nodes. There are $K$ modality planes, each with
adjacency matrices $\mathbf{A}_{(k)}\in\mathcal{R}^{P\times P}$. The $K\times
K$ pairwise cross planar connections are given by
$\mathbf{C}_{(l,m)}\in\mathcal{R}^{P\times P}$, where $P=|{\mathcal{V}}|$. All
edge weights may take values in range $[0-1]$

###  2.2 HGR Maximal Correlations for Latent Multi-Graph Learning:

Recall that we would like to learn task informed patient-modality multi-graph
representations automatically from unstructured modality data. To this end, we
develop the framework illustrated in the Gray Box in Fig. 1.

Let $\\{\mathbf{x}_{n}^{k}\in\mathcal{R}^{D_{k}\times 1}\\}^{K}_{k=1}$ be the
features from to modality $k$ for patient $n$. Since features from different
modalities lie in different input subspaces, we develop parallel common space
projections to explore the dependence between them. The Hirschfield, Gebelin,
Rènyi (HGR) [23] framework in statistics is known to generalize the notion of
dependence to abstract and non-linear functional spaces. Such non-linear
projections can be parameterized by deep neural networks.

Specifically, let the collection of modality-specific projection networks be
given by $\\{\mathbf{f}^{k}(\cdot):\mathcal{R}^{D_{k}\times
1}\rightarrow\mathcal{R}^{D_{p}\times 1}\\}$. The HGR maximal correlation is a
symmetric measure obtained by solving the following coupled pairwise
constrained optimization problem:

$\sup_{\mathcal{C}_{E},\mathcal{C}_{\text{Cov}}}{\rho_{\text{HGR}}(\mathbf{x}^{l},\mathbf{x}^{m})}=\sup_{\mathcal{C}_{E},\mathcal{C}_{\text{Cov}}}{\mathbb{E}\Big{[}[\mathbf{f}^{l}(\mathbf{x}^{l})]^{T}\mathbf{f}^{m}(\mathbf{x}^{m})\Big{]}}$
(1)

$\forall\\{l.m\\}\ \text{s.t.}\ l\neq m$, where $\mathbf{\mathcal{I}}_{D_{p}}$
is a $D_{p}\times D_{p}$ identity matrix. The constraint sets are given by:

$\displaystyle\mathcal{C}_{E}:\\{\mathbb{E}[\mathbf{f}^{l}(\cdot)]=\mathbb{E}[\mathbf{f}^{m}(\cdot)]=\mathbf{0}\\}\
\ \ \ $ (2)
$\displaystyle\mathcal{C}_{\text{Cov}}:\\{\mathbf{Cov}(\mathbf{f}^{l}(\mathbf{x}^{l}))=\mathbf{Cov}(\mathbf{f}^{m}(\mathbf{x}^{m}))=\mathbf{\mathcal{I}}_{D_{p}}\\}$
(3)

Approaches such as deep CCA [25] can be thought of as a special case of this
formulation which solve the whitening (empricial covariance) constraints (Eq.
(3)) via explicit pairwise de-correlation.

However, for multiple modalities in large datasets, exact whitening is not
scalable. To circumvent this issue, we can use the approach in [23]. This
formulation proposes introduces a relaxation to the exact HGR, named soft-HGR,
which consists of a trace regularizer in lieu of whitening. Eq. (1) can be
relaxed as an empirical minimization problem $\text{min}\
\mathcal{L}_{\text{sHGR}}$, where the sHGR loss is:

$\displaystyle\mathcal{L}_{\text{sHGR}}=-\frac{1}{N_{z}}\mathbb{E}\Big{[}\mathbf{f}^{l}(\mathbf{x}^{l})^{T}\mathbf{f}^{m}(\mathbf{x}^{m})\Big{]}+\frac{1}{2N_{z}}\operatorname{Tr}\Big{[}{\mathbf{Cov}[\mathbf{f}^{l}(\mathbf{x}^{l})]}{\mathbf{Cov}[\mathbf{f}^{m}(\mathbf{x}^{m})]}\Big{]}$
(4)
$\displaystyle=-\frac{1}{N_{z}}\Bigg{(}\sum^{K}_{l,m=1}\Bigg{[}\sum^{N}_{n=1}\frac{\mathbf{f}^{l}(\mathbf{x}_{n}^{l})^{T}\mathbf{f}^{m}(\mathbf{x}_{n}^{m})}{(N-1)}-\frac{\operatorname{Tr}\Big{[}{\mathbf{Cov}[\mathbf{f}^{l}(\mathbf{x}^{l})]}{\mathbf{Cov}[\mathbf{f}^{m}(\mathbf{x}^{m})]}\Big{]}}{2}\Bigg{]}\Bigg{)}\
\ l\neq m$

The expectation under the functional transformations
$\mathbb{E}[\mathbf{f}^{l}(\mathbf{x}^{l})]=\mathbb{E}[\mathbf{f}^{m}(\mathbf{x}^{m})]=\mathbf{0}$
is enforced step-wise by mean subtraction during optimization. Here,
$\mathbf{Cov}(\cdot)$ is the empirical covariance matrix. We parameterize
$\\{\mathbf{f}^{k}(\cdot)\\}$ as a simple two layered fully connected neural
network with a normalization factor as $N_{z}=M(M-1)$.

By design, the MaxCorr formulation allows us to utilize the correlation
$\rho_{\text{sHGR}}(\mathbf{x}_{i}^{l},\mathbf{x}_{j}^{m})$ (computed after
solving Eq. (4)) to model dependence between patients $i$ and $j$ according to
the $l$ and $m$ modality features in a general setting. The absolute value of
this correlation measure define the edge weights between nodes in the patient-
modality multi-graph. As opposed to existing

#### Learnable Adaptive Sparsity:

Additionally, we would like to have our learning framework automatically
discover and retain salient edges that are relevant for prediction. To
encourage sparsity in the edges, we utilize a learnable soft-thresholding
formulation. We first define a symmetric block sparsity matrix $\mathbf{S}$.
Since edge weights in the multi-graph are in the range $[0-1]$, we normalize
it through the sigmoid function as
$\tilde{\mathbf{S}}=\tilde{\mathbf{S}}^{T}=\text{Sigmoid}({\mathbf{S}})\in\mathcal{R}^{K\times
K}$. The entries of the soft-thresholding matrix $\tilde{\mathbf{S}}[l,m]$
define learnable thresholds for the cross modal connections when $l\neq m$ and
in-plane connections when $l=m$. Finally, the cross modal edges and in-plane
edges of the multi-graph are given by

$\displaystyle\mathbf{C}_{(l,m)}[i,j]=\text{ReLU}({\tilde{\rho}_{\text{sHGR}}(\mathbf{x}_{i}^{l},\mathbf{x}_{j}^{m})}-\tilde{\mathbf{S}}[l,m])$
(5)
$\displaystyle\mathbf{A}_{k}[i,j]=\text{ReLU}(\tilde{\rho}_{\text{sHGR}}(\mathbf{x}_{i}^{k},\mathbf{x}_{j}^{k})-\tilde{\mathbf{S}}[k,k])$
(6)

with ${\tilde{\rho}_{\text{sHGR}}}=|{\rho}_{\text{sHGR}}|$ respectively. The
adjacency matrices $\mathbf{A}_{(k)}$ model the dependence within the features
of modality $k$, while the cross planar matrices $\\{\mathbf{C}_{(l,m)}\\}$
capture interactions across modalities. Overall, $\mathbf{S}$ acts as a
regularizer that suppresses noisy weak dependencies. These regularization
parameters are automatically inferred during training along with the MaxCorr
projection parameters $\\{\mathbf{f}^{k}(\cdot)\\}$. This effectively adds
just $K(K+1)$ learnable parameters to the MaxCorrMGNN.

###  2.3 Multi-Graph Neural Network:

As a standalone optimization, the MaxCorr block is not guaranteed to learn
discriminative projections of modality features. A natural next step is to
couple the multi-graph representation learning with the classification task.
Graph Neural Networks have recently become popular tools for reasoning from
graphs and graph-signals. Given the patient-modality multi-graph, we design an
extension of traditional graph neural networks to multi-graphs for inference
tasks.

Conventional GNNs filter information based on the graph topology (i.e. the
adjacency matrix) to map the node features to the targets based on graph
traversals. Conceptually, cascading $l$ GNN layers is analogous to filtered
information pooling at each node from its $l$-hop neighbors [13] inferred from
the powers of the graph adjacency matrix. These neighborhoods can be reached
by seeding walks on the graph starting at the desired node. Inspired by this
design, we craft a multi-graph neural network (Purple Box in Fig. 1) for
outcome prediction. Our MGNN generalizes structured message passing to the
multi-graph in a manner similar to those done for multiplexed graphs [6].
Notably, our formulation is more general, as it avoids using strictly vertical
interaction constraints between patients across modalities.

We first construct two supra-adjacency matrices to perform walks on the multi-
graph $\mathcal{G}_{\text{M}}$ for fusion. The first is the intra-modality
adjacency matrix $\boldsymbol{\mathcal{A}}\in\mathcal{R}^{PK\times PK}$. The
second is the inter-modality connectivity matrix
${\boldsymbol{\mathcal{C}}}\in\mathcal{R}^{PK\times PK}$, each defined block-
wise. Mathematically, we express this as:

$\displaystyle\boldsymbol{\mathcal{A}}=\bigoplus_{k}\mathbf{A}_{(k)}$ (7)
$\displaystyle\hat{\boldsymbol{\mathcal{C}}}:\hat{\boldsymbol{\mathcal{C}}}[lP:(l+1)P,mP:(m+1)P]=\mathbf{C}_{(l,m)}\
\mathbb{1}(l\neq m)+\boldsymbol{\mathcal{I}_{P}}\ \mathbb{1}(l=m)$ (8)

where $\bigoplus$ is the direct sum operation and $\mathbb{1}$ denotes the
indicator function. By design, $\boldsymbol{\mathcal{A}}$ is block-diagonal
and allows for within-planar (intra-modality) transitions between nodes. The
off-diagonal blocks of $\hat{\boldsymbol{\mathcal{C}}}$, i.e.
$\mathbf{C}_{(l,m)}$, capture transitions between nodes as per cross-planar
(inter-modality) relationships.

#### MGNN Message Passing Walks:

Walks on $\mathcal{G}_{\text{M}}$ combine within and across planar steps to
reach a patient supra-node $s_{j}$ from another supra-node $s_{i}$
($s_{i},s_{j}\in\mathcal{V}_{\text{M}}$). We characterize the multi-hop
neighborhoods and transitions using factorized operations involving
$\boldsymbol{\mathcal{A}}$ and $\hat{\boldsymbol{\mathcal{C}}}$. We can
perform a multi-graph walk via two types of distinct steps, i.e., (1) an
isolated intra-planar transition or (2) a transition involving an inter-planar
step either before or after a within-planar step. These steps can be
exhaustively recreated via two factorizations: (I) after one intra-planar
step, the walk may continue in the same modal plane or hop to a different one
via $\boldsymbol{\mathcal{A}}\hat{\boldsymbol{\mathcal{C}}}$ and (II) the walk
may continue in the current modal plane or hop to a different plane before the
intra-planar step via
$\hat{\boldsymbol{\mathcal{C}}}\boldsymbol{\mathcal{A}}$.

The Multi-Graph Neural Network (MGNN) uses these walk operations to
automatically mine predictive patterns from the multi-graph given the targets
(task-supervision) and the GNN parameters. For supra-node $s_{i}$,
$\mathbf{h}^{(d)}_{s_{i}}\in\mathcal{R}^{D^{d}\times 1}$ is the feature
(supra)-embedding at MGNN depth $d$. The forward pass operations of the MGNN
are as follows:

$\displaystyle\mathbf{h}_{s_{i},I}^{(d+1)}=\boldsymbol{\phi}^{(d)}_{I}\Big{(}(1+\epsilon)\mathbf{h}^{(d)}_{s_{i}}+\text{wmean}\Big{[}\mathbf{h}^{(d)}_{s_{j}},\boldsymbol{\mathcal{A}}\mathbf{\hat{\boldsymbol{\mathcal{C}}}}[s_{i},s_{j}]\
\ ;\ \
s_{j}\in\mathcal{N}_{\boldsymbol{\mathcal{A}}\hat{\boldsymbol{\mathcal{C}}}}(s_{i})\Big{]}\Big{)}\
\ \ \ \ \ $ (9)
$\displaystyle\mathbf{h}_{s_{i},II}^{(d+1)}=\boldsymbol{\phi}^{(d)}_{II}\Big{(}(1+\epsilon)\mathbf{h}^{(d)}_{s_{i}}+\text{wmean}\Big{[}\mathbf{h}^{(d)}_{s_{j}},\mathbf{\hat{\boldsymbol{\mathcal{C}}}\boldsymbol{\mathcal{A}}}[s_{i},s_{j}]\
\ ;\ \
s_{j}\in\mathcal{N}_{\hat{\boldsymbol{\mathcal{C}}}\boldsymbol{\mathcal{A}}}(s_{i})\Big{]}\Big{)}\
\ \ \ \ \ $ (10)
$\displaystyle\mathbf{h}^{(d+1)}_{s_{i}}=\text{concat}(\mathbf{h}^{(d+1)}_{s_{i},I},\mathbf{h}^{(d+1)}_{s_{i},II})\
\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ $ (11) $\displaystyle\ \ \ \ \ \ \
\ \ \ \ \ \ \
\mathbf{g}_{o}(\\{\mathbf{h}^{(L)}_{s_{i}}\\}_{s_{i}\leftrightarrow
i})=\hat{\mathbf{Y}}_{i}\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \
$ (12)

At the input layer, we have
$\mathbf{h}_{s_{i}}^{(0)}=\mathbf{f}^{k}(\mathbf{x}_{i}^{k})$ computed from
the modality features for patient $i$ after the sHGR transformation from the
corresponding modality $k$. We then concatenate the supra-embeddings as input
to the next layer i.e. $\mathbf{h}^{(d+1)}_{s_{i}}$. Eqs. (9-10) denote the
Graph Isomorphism Network (GIN) [24] with
$\\{\boldsymbol{\mathcal{\phi}}^{(d)}_{I}(\cdot),\boldsymbol{\mathcal{\phi}}^{(d)}_{II}(\cdot)\\}$
as layerwise linear transformations. This performs message passing on the
multi-graph using the neighborhood relationships and normalized edge weights
from the walk matrices in the weighted mean operation $\text{wmean}(\cdot)$.

From the interpretability standpoint, these operations keep the semantics of
the embeddings intact at both the patient and modality level throughout the
MaxCorrMGNN transformations. Finally, $\mathbf{g}_{o}(\cdot)$ is a graph
readout network that maps to the one-hot encoded outcome $\mathbf{Y}$, which
performs a convex combination of the filtered modality embeddings, followed by
a linear readout.

###  2.4 End-to-end Learning through task supervision:

Piecing together the constituent components, i.e. the latent graph learning
and MGNN inference module, we optimize the following coupled objective
function:

$\mathcal{L}=\lambda\mathcal{L}_{\text{sHGR}}+(1-\lambda)\mathcal{L}_{\text{CE}}(\hat{\mathbf{Y}},\mathbf{Y})$
(13)

with $\lambda\in[0,1]$ being a tradeoff parameter and
$\mathcal{L}_{CE}(\cdot)$ being the cross entropy loss. The parameters
$\\{\\{\mathbf{f}^{k}(\cdot)\\},\mathbf{S},\\{\boldsymbol{\mathcal{\phi}}^{(d)}_{I}(\cdot),\boldsymbol{\mathcal{\phi}}^{(d)}_{II}(\cdot),\epsilon\\},\mathbf{g}_{o}(\cdot)\\}$
of the framework are jointly learned via standard backpropagation.

#### Inductive Learning for Multi-Graph Neural Networks:

The multi-graph is designed to have subjects as the nodes, which requires us
to adapt training to accomodate an inductive learning setup. Specifically, we
train the MaxCorrMGNN in a fully supervised fashion by extending the
principles outlined in [2] for multi-layered graphs. During training, we use
only the supra-node features and induced sub-graph edges (including both
cross-modal and intra-planar edges) associated with the subjects in the
training set for backpropagation. During validation/testing, we freeze the
parameter estimates and add in the edges corresponding to the unseen patients
to perform a forward pass for estimation. This procedure ensures that no
double dipping occurs in the hyper-parameter estimation, nor in the evaluation
step. Additionally, while not the focus of this application, this procedure
allows for extending prediction and training to an online setting, where new
subject/modality information may dynamically become available.

#### Implementation Details:

We implement the MaxCorr projection networks $\mathbf{f}^{l}(\cdot)$ as a
simple three layered neural network with hidden layer width of $32$ and output
$D_{p}=64$ and LeakyReLU activation (negative slope=0.01). The MGNN layers are
Graph Isomorphism Network (GIN)[24] with ReLU activation and linear readout
(width:$64$) and batch normalization. $\mathbf{g}_{o}(\cdot)$ implements a
convex combination of the modality embeddings followed by a linear layer. We
use the ADAMw optimizer [15] and train on a 64GB CPU RAM, 2.3 GHz 16-Core
Intel i9 machine (18-20 min training time per run). We set the hyperparameters
for our model (and baselines) using grid-search to $\lambda=0.01$, learning
rate$=0.0001$, weight decay$=0.001$, epochs$=50$, batch size$=128$ after pre-
training the network on the sHGR loss alone for $50$ epochs. All frameworks
are implemented on the Deep Graph Library (v=0.6.2) in PyTorch (v=0.10.1).

##  3 Experiments and Results

###  3.1 Data and Preprocessing

Figure 2: NIH TB Dataset: Multimodal data for treatment outcome prediction.

We evaluate our model on the Tuberculosis Data Exploration Portal [10]
consisting of 3051 patients with five different treatment outcomes (Died,
Still on treatment, Completed, Cured, or Failure) with the class frequencies
as: 0.21/0.11/0.50/0.10/0.08 respectively and five modalities. We pre-process
the data according to the procedure outlined in [6].

For each subject, we have features available from demographic, clinical,
regimen and genomic recordings with chest CTs available for 1015 of them. We
have a total of 4081 genomic, 29 demographic, 1726 clinical, 233 regimen
features that are categorical, and 2048 imaging and 8 miscellaneous continuous
features. Information that may directly be related to treatment outcomes, eg
drug resistance type, were removed from the clinical and regimen features.

For genomic data, 81 single nucleotide polymorphisms (SNPs) from the causative
organisms Mycobacterium tuberculosis (Mtb) known to be related to drug
resistance were used. For 275 of the subjects, we also assemble the raw genome
sequence from NCBI Sequence Read Archive. This provides a more fine-grained
description of the biological sequences of the causative pathogen [17].
Briefly, we performed a de novo assembly process on each Mtb genome to yield
protein and gene sequences. We utilized InterProScan [12] to further process
the protein sequences and extract the functional domains, i.e. sub-sequences
located within the protein’s amino acid chain responsible for the enzymatic
bioactivity of a protein. This provides a total of 4000 functional genomic
features. Finally, for the imaging modality, the lung was segmented via multi-
atlas segmentation [22] followed by a pre-trained DenseNet [11] to extract a
1024-dimensional feature vector for each axial slice intersecting the lung.
The mean and maximum of each feature were then assembled to give a total of
2048 features. Missing features are imputed from the training cohort using
mean imputation for all runs.

###  3.2 Evaluation Metrics:

Since we have a five-class classification task, we evaluate the prediction
performance of the MaxCorrMGNN and the baselines using the AU-ROC (Area Under
the Receiver Operating Curve) metric. Given the prediction logits, this metric
is computed both class-wise and as a weighted average. Higher per-class and
overall AU-ROC indicate improved performance. For our experiments, we use $10$
randomly generated train/validation/test splits with ratio 0.7/0.1/0.2 to
train our model and each baseline.

Finally, statistical differences between the baselines and our method are
measured according to the DeLong [4] test computed class-wise. This test is a
sanity check to evaluate whether perceived differences in model performance
are robust to sampling.

###  3.3 Baseline Comparisons

We perform a comprehensive evaluation of our framework for the problem of
multimodal fusion. Our baseline comparisons can be grouped into three
categories, namely, (1) Single Modality Predictors/ No Fusion (2) State-of-
the-art Conventional including early/late/intermediate fusion and Latent-Graph
Learning models from literature (3) Ablation Studies.

The ablation studies evaluate the efficacy of the three main constituents of
the MaxCorrMGNN, i.e. the MaxCorr graph construction, the Multi-Graph Neural
Network and the end-to-end optimization.

  * •

Single Modality: For this comparison, we run predictive deep-learning models
on the individual modality features without fusing them as a benchmark. We use
a two layered multi-layered perception (MLP) with hidden layer widths as 400
and 20 and LeakyReLU activation (neg. slope=0.01).

  * •

Early Fusion: For early fusion, individual modality features are first
concatenated and then fed through a neural network. The predictive model has
the same architecture as the previous baseline.

  * •

Uncertainty Based Late Fusion [21]: We combine the predictions from the
individual modalities in the previous baseline using a state-of-the-art late
fusion framework in [21]. This model estimates the uncertainty in the
individual classifiers to improve the robustness of outcome prediction. Unlike
our work, patient-modality dependence is not explicitly modeled as the
modality predictions are only combined after individual modality-specific
models have been trained. Hyperparameters are set according to [21].

  * •

Graph Based Intermediate Fusion [6]: This is a graph based neural framework
that achieved state-of-the-art performance on multimodal fusion on
unstructured data. This model follows a two step procedure. For each patient,
this model first converts the multimodal features into a fused binary
multiplex graph (multi-graph where all blocks of
$\mathbf{\hat{\boldsymbol{\mathcal{C}}}}$ are strictly diagonal) between
features. The graph connectivity is learned in an unsupervised fashion through
auto-encoders. Following this, a multiplexed graph neural network is used for
inference. Hyperparameters are set according to [6]. While this framework
takes a graph based approach to fusion, the construction of the graph is not
directly coupled with the task supervision.

  * •

Latent Graph Learning [2]: This baseline was developed for fusing multimodal
data for prediction. It introduces a latent patient-patient graph learning
from the concatenated modality features via a graph-attention (GAT-like [20])
formulation. However, unlike our model, the feature concatenation does not
distinguish between intra- and inter-modality dependence across patients i.e.
it constructs a single-relational (collapsed) graph that is learned as a part
of the training.

  * •

sHGR+ANN [23]: This is a state-of-the-art multimodal fusion framework [23]
that also utilizes the sHGR formulation to infer multi-modal data
representations. However, instead of constructing a patient-modality graph,
the projected features are combined via concatenation. Then, a two layered MLP
(hidden size:200) maps to the outcomes, with the two objectives trained end-
to-end. This baseline can be thought of as an ablation that evaluates the
benefit of using the multi-graph neural network for fine-grained reasoning.
Additionally, this and the previous framework help us evaluate the benefit of
our patient-modality multi-graph representation for fusion.

  * •

MaxCorrMGNN w/o sHGR: Through this comparison, we evaluate the need for using
the soft HGR formulation to construct the latent multi-graph. Keeping the
architectural components consistent with our model, we set $\lambda=0$ in Eq.
(13). Note that this ablation effectively converts the multi-graph
representation learning into a modality specific self/cross attention
learning, akin to graph transformers. Overall, this framework helps us
evaluate the benefit of our MaxCorr formulation for latent multi-graph
learning.

  * •

Decoupled MaxCorrMGNN: Finally, this ablation is designed to examine the
benefit of coupling the MaxCorr and MGNN into a coupled objective. Therefore,
instead of an end-to-end training, we run the sHGR optimization first,
followed by the MGNN for prediction.

Figure 3: We display the mean per-class and weighted average AU-ROC and the
standard errors for TB outcome prediction against (Left): Single Modality
Predictors (Middle): Traditional and Graph Based Fusion Frameworks (Right):
Ablations of the MaxCorrMGNN. * indicate comparisons against the MaxCorrMGNN
according to the DeLong test that achieve statistical significance ($p<0.01$).

###  3.4 Outcome Prediction Performance:

Fig. 3 illustrates the outcome prediction performance of our framework against
the single modality predictors (left), state-of-the-art fusion frameworks
(middle), and ablated versions of our model (right). Comparisons marked with
$*$ achieve a statistical significance threshold of $p<0.01$ across runs as
per the DeLong test [4]. Note that our fusion framework outperforms all of the
single modality predictors by a large margin. Moreover, the traditional and
graph-based fusion baselines also provide improved performance against the
single modality predictors. Taken together, these observations highlight the
need for fusion of multiple modalities for outcome prediction in TB. This
observation is consistent with findings in treatment outcome prediction
literature [6, 16] in TB.

The MaxCorrMGNN also provides improved performance when compared to all of the
fusion baselines, with most comparisons achieving statistical significance
thresholds. While the Early Fusion and Uncertainty based Late fusion [21]
networks provide marked improvements over single modality predictions, but
still fail to reach the performance level of our model. This is likely due to
their limited ability to leverage subtle patient-specific cross-modal
interactions.

On the other hand, the latent graph learning in [2] models connectivity
between subjects as a part of the supervision. However, this method collapses
the different types of dependence into one relation-type, which may be too
restrictive for fusion applications. The intermediate fusion framework of [6]
was designed to address these limitations by the use of multiplex graphs.
However, the artificial separation between the graph construction and
inference steps may not inherently extract discriminative multi-graph
representations, which could explain the performance gap against our
framework.

Finally, the three ablations, the sHGR+ANN [23], MaxCorrMGNN w/o sHGR,
Decoupled MaxCorrMGNN help us systematically examine the three building blocks
of our framework, i.e. MGNN and MaxCorr networks individually as well as the
end-to-end training of the two blocks. We observe a notable performance drop
in these baselines, which reinforces the principles we considered in carefully
designing the individual components. In fact, the comparison against the
Decoupled MaxCorrMGNN illustrates that coupling the two components into a
single objective is key to obtaining improved representational power for
predictive tasks. Taken together, our results suggest that the MaxCorrMGNN is
a powerful framework for multimodal fusion of unstructured data.

##  4 Discussion

We have developed a novel multi-graph deep learning framework, i.e. the
MaxCorrMGNN for generalized problems of multimodal fusion in medical data.
Going one step beyond simple statistical measures, the patient-modality multi-
layered graph allows us to uncover nuanced non-linear notions of dependence
between modality features via the maximal correlation soft-HGR formulation.
The sHGR formulation coupled with the learnable sparsity module allow us to
directly translate an abstract measure of interaction across subjects and
modalities in any multimodal dataset into a patient-modality multi-layered
graph structure for inference. The construction of the multi-graph planes
allow the node features to retain their individuality in terms of the plane
(modality) and patient (node-identity) in the filtered Graph Neural Network
representations. This admits more explainable intermediate representations in
comparison to the baselines, i.e. provides us with the ability to explicitly
reason at the granularity of both the subjects and modalities. Conversely, the
graph based/traditional fusion baselines collapse this information, either in
the multimodal representation or in the inference step. We believe that this
added flexibility in the MaxCorrMGNN contributes to the improved
generalization power in practice. Finally, all the individual components (i.e.
MaxCorr, learnable soft-thresholding, MGNN message passing) are designed to be
fully differentiable deep learning operations, allowing us to directly couple
them end-to-end. We demonstrate in experiment that this coupling is key to
generalization. As such, this model makes very mild assumptions about the
nature of the multimodal data. The general principles and machinery developed
in this work would likely be useful to a wide variety of applications beyond
the medical realm.

#### Limitations and Future Work:

In problems of multimodal fusion, especially for medical applications, data
acquisition is a fairly contrived and expensive process. In many real-world
modalities may often be only partially observed, missing in totality, or noisy
in acquisition. Simple methods such as mean based imputation may be inadequate
for fine-grained reasoning. As an aim to address this, an active line of
exploration is to extend the framework to handle missing, ambiguous and
erroneous data and labels within the multilayered graph representation. This
may be achieved by leveraging statistical and graph theoretic tools that can
be integrated directly into the message passing walks. Finally, the multi-
graph and HGR construction focuses on uncovering pairwise relationships
between subjects and features. A future direction would be to extend these
frameworks to model complex multi-set dependencies.

##  5 Conclusion

We have introduced a novel multi-layered graph based neural framework for
general inference problems in multimodal fusion. Our framework leverages the
HGR MaxCorr formulation to convert unstructured multi-modal data into a
patient-modality multi-graph. We design a generalized multi-graph neural
network for fine-grained reasoning from this representation. Our design
preserves the patient-modality semantics as a part of the architecture, making
our representations more readily interpretable rather than fully black-box.
The end-to-end optimization of the two components offers a viable tradeoff
between flexibility, representational power, and interpretability. We
demonstrate the efficacy of the MaxCorr MGNN for fusing disparate information
from imaging, genomic and clinical data for outcome prediction in Tuberculosis
and demonstrate consistent improvements against competing state-of-the-art
baselines developed in literature. Moreover, the framework makes very few
assumptions making it potentially applicable to a variety of fusion problems
in other AI domains. Finally, the principles developed in this paper are
general and can potentially be applied to problems well beyond multimodal
fusion.

## References

  * [1] Baltrušaitis, T., Ahuja, C., Morency, L.P.: Multimodal machine learning: A survey and taxonomy. IEEE transactions on pattern analysis and machine intelligence 41(2), 423–443 (2018)
  * [2] Cosmo, L., Kazi, A., Ahmadi, S.A., Navab, N., Bronstein, M.: Latent-graph learning for disease prediction. In: Medical Image Computing and Computer Assisted Intervention–MICCAI 2020: 23rd International Conference, Lima, Peru, October 4–8, 2020, Proceedings, Part II 23. pp. 643–653. Springer (2020)
  * [3] Cozzo, E., de Arruda, G.F., Rodrigues, F.A., Moreno, Y.: Multiplex networks (2018). https://doi.org/10.1007/978-3-319-92255-3, http://link.springer.com/10.1007/978-3-319-92255-3
  * [4] DeLong, E.R., DeLong, D.M., Clarke-Pearson, D.L.: Comparing the areas under two or more correlated receiver operating characteristic curves: a nonparametric approach. Biometrics pp. 837–845 (1988)
  * [5] Dsouza, N.S., Nebel, M.B., Crocetti, D., Robinson, J., Mostofsky, S., Venkataraman, A.: M-gcn: A multimodal graph convolutional network to integrate functional and structural connectomics data to predict multidimensional phenotypic characterizations. In: Medical Imaging with Deep Learning. pp. 119–130. PMLR (2021)
  * [6] D’Souza, N.S., Wang, H., Giovannini, A., Foncubierta-Rodriguez, A., Beck, K.L., Boyko, O., Syeda-Mahmood, T.: Fusing modalities by multiplexed graph neural networks for outcome prediction in tuberculosis. In: Medical Image Computing and Computer Assisted Intervention–MICCAI 2022: 25th International Conference, Singapore, September 18–22, 2022, Proceedings, Part VII. pp. 287–297. Springer (2022)
  * [7] D’Souza, N.S., Nebel, M.B., Crocetti, D., Robinson, J., Mostofsky, S., Venkataraman, A.: A matrix autoencoder framework to align the functional and structural connectivity manifolds as guided by behavioral phenotypes. In: Medical Image Computing and Computer Assisted Intervention–MICCAI 2021: 24th International Conference, Strasbourg, France, September 27–October 1, 2021, Proceedings, Part VII 24. pp. 625–636. Springer (2021)
  * [8] D’Souza, N.S., Nebel, M.B., Crocetti, D., Robinson, J., Wymbs, N., Mostofsky, S.H., Venkataraman, A.: Deep sr-ddl: Deep structurally regularized dynamic dictionary learning to integrate multimodal and dynamic functional connectomics data for multidimensional clinical characterizations. NeuroImage 241, 118388 (2021)
  * [9] D’Souza, N.S., Nebel, M.B., Crocetti, D., Wymbs, N., Robinson, J., Mostofsky, S., Venkataraman, A.: A deep-generative hybrid model to integrate multimodal and dynamic connectivity for predicting spectrum-level deficits in autism. In: Medical Image Computing and Computer Assisted Intervention–MICCAI 2020: 23rd International Conference, Lima, Peru, October 4–8, 2020, Proceedings, Part VII 23. pp. 437–447. Springer (2020)
  * [10] Gabrielian, A., Engle, E., Harris, M., Wollenberg, K., Juarez-Espinosa, O., Glogowski, A., Long, A., Patti, L., Hurt, D.E., Rosenthal, A., Tartakovsky, M.: Tb depot (data exploration portal): A multi-domain tuberculosis data analysis resource. PLOS ONE 14(5), e0217410 (may 2019). https://doi.org/10.1371/journal.pone.0217410, http://dx.plos.org/10.1371/journal.pone.0217410
  * [11] Huang, G., Liu, Z., Van Der Maaten, L., Weinberger, K.Q.: Densely connected convolutional networks. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 4700–4708 (2017)
  * [12] Jones, P., Binns, D., Chang, H.Y., Fraser, M., Li, W., McAnulla, C., McWilliam, H., Maslen, J., Mitchell, A., Nuka, G., Pesseat, S., Quinn, A.F., Sangrador-Vegas, A., Scheremetjew, M., Yong, S.Y., Lopez, R., Hunter, S.: InterProScan 5: genome-scale protein function classification. Bioinformatics (Oxford, England) 30(9), 1236–40 (may 2014). https://doi.org/10.1093/bioinformatics/btu031
  * [13] Kipf, T.N., Welling, M.: Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907 (2016)
  * [14] Lahat, D., Adali, T., Jutten, C.: Multimodal data fusion: an overview of methods, challenges, and prospects. Proceedings of the IEEE 103(9), 1449–1477 (2015)
  * [15] Loshchilov, I., Hutter, F.: Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 (2017)
  * [16] Muñoz-Sellart, M., Cuevas, L., Tumato, M., Merid, Y., Yassin, M.: Factors associated with poor tuberculosis treatment outcome in the southern region of ethiopia. The International Journal of Tuberculosis and Lung Disease 14(8), 973–979 (2010)
  * [17] Seabolt, E.E., Nayar, G., Krishnareddy, H., Agarwal, A., Beck, K.L., Terrizzano, I., Kandogan, E., Roth, M., Mukherjee, V., Kaufman, J.H.: OMXWare, A Cloud-Based Platform for Studying Microbial Life at Scale (nov 2019), http://arxiv.org/abs/1911.02095
  * [18] Subramanian, V., Do, M.N., Syeda-Mahmood, T.: Multimodal fusion of imaging and genomics for lung cancer recurrence prediction. In: 2020 IEEE 17th International Symposium on Biomedical Imaging (ISBI). pp. 804–808. IEEE (2020)
  * [19] Subramanian, V., Syeda-Mahmood, T., Do, M.N.: Multi-modality fusion using canonical correlation analysis methods: Application in breast cancer survival prediction from histology and genomics. arXiv preprint arXiv:2111.13987 (2021)
  * [20] Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., Bengio, Y.: Graph attention networks. arXiv preprint arXiv:1710.10903 (2017)
  * [21] Wang, H., Subramanian, V., Syeda-Mahmood, T.: Modeling uncertainty in multi-modal fusion for lung cancer survival analysis. In: 2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI). pp. 1169–1172. IEEE (2021)
  * [22] Wang, H., Yushkevich, P.: Multi-atlas segmentation with joint label fusion and corrective learning—an open source implementation. Frontiers in neuroinformatics 7, 27 (2013)
  * [23] Wang, L., Wu, J., Huang, S.L., Zheng, L., Xu, X., Zhang, L., Huang, J.: An efficient approach to informative feature extraction from multimodal data. In: Proceedings of the AAAI Conference on Artificial Intelligence. vol. 33, pp. 5281–5288 (2019)
  * [24] Xu, K., Hu, W., Leskovec, J., Jegelka, S.: How powerful are graph neural networks? arXiv preprint arXiv:1810.00826 (2018)
  * [25] Yang, X., Liu, W., Liu, W., Tao, D.: A survey on canonical correlation analysis. IEEE Transactions on Knowledge and Data Engineering 33(6), 2349–2368 (2019)
  * [26] Zheng, S., Zhu, Z., Liu, Z., Guo, Z., Liu, Y., Yang, Y., Zhao, Y.: Multi-modal graph learning for disease prediction. IEEE Transactions on Medical Imaging 41(9), 2207–2216 (2022)