Bayesian Relational Generative Model for Scalable Multi-modal Learning

1 INTRODUCTION . Many prominent methods for multi-modal learning ( Argelaguet et al. , 2018 ; Andrew et al. , 2013 ; Klami et al. , 2013 ; Zhao et al. , 2016 ) are based on the canonical correlation analysis ( CCA ) , which extracts shared components across multiple views . The main idea is that , given two vectors of random variables , the method finds the linear projections in a shared latent space , in which the projected vectors are maximally correlated ( Thompson , 1984 ) . This can help to understand the overall dependency structure between these two vectors . However , the classical CCA can not handle non-linearity ( Andrew et al. , 2013 ) and suffers from a lack of probabilistic interpretation when applied to high dimensional data ( Klami et al. , 2013 ) . To address these issues , probabilistic CCA ( PCCA ) has been proposed and extended to non-linear settings using kernel methods and neural networks ( Bach & Jordan , 2005 ) . More recently , Wu & Goodman ( 2018 ) introduced the multimodal variational autoencoder ( mVAE ) that models the joint posterior as a product-of-experts ( PoE ) over the marginal posteriors , enabling cross-modal generation at test time without requiring additional inference networks and multi-stage training regimes . However , training a PoE is difficult and the following techniques are needed in order to ensure that the individual views are learnt faithfully ; i ) using artificial sub-sampling of the observed views ( Wu & Goodman , 2018 ) , ii ) : applying variants of contrastive divergence ( Hinton , 2002 ) , and iii ) : utilizing information bottleneck on the marginal representations of each view ( Lee & Schaar , 2021 ) . Furthermore , each expert holds the power of veto and low density of only one marginal posteriors among a given set of observations leads to the low density of joint distribution ( Shi et al. , 2019 ) . In the case of Gaussian experts , different levels of complexity of modalities or sub-optimal initialisation might result in the miscalibrated precisions that potentially lead to a biased overall mean prediction ( Shi et al. , 2019 ) . Despite the success of the existing multi-modal generative methods , they suffer from inducing stochasticity in their inferred functions . This leads to ; i ) poor adaptation to new observations1 , ii ) inefficient generalizability in few-shot settings and , iii ) inaccurate uncertainty estimation , all of which are important in many safety critical applications such as healthcare . On the other hand , stochastic processes , such as Gaussian processes ( GPs ) , are exchangeable models that posit distributions over possible functions , and are updated in light of data through the probabilistic inference . Despite these advantages , GPs are computationally expensive and their underlying model is not flexible for high-dimensional inputs , making them infeasible for multi-modal learning settings with multiple heterogeneous and high-dimensional data types . Apart from the aforementioned issues , existing generative models for multi-modal learning focus on latent representation , but do not fully incorporate the label information . In order to facilitate predictions on the target labels , they need to apply two-step procedures , in which the shared latent spaces are used for downstream classification . While information relevant to the reconstruction of views will be well-captured in the shared representations , discriminative information relevant to the target task may be discarded , leading to poor prediction performance of such models ( Lee & Schaar , 2021 ) . Furthermore , the shared representations of these methods in high-dimensional space are often difficult to interpret and need an extra downstream analyses such as PCA or t-SNE ( Van Der Maaten , 2014 ) that may cause biased results ( Lötsch & Ultsch , 2020 ) . In this paper , instead of following the existing multi-modal generative methods that rely on deterministic predictive functions , we propose the first multi-modal stochastic processes family that learns distributions over functions for inputs with any missing patterns . Like GPs , it naturally provides robust uncertainty estimates and can encode inductive biases . Unlike GPs , though , mRNP is computationally efficient during training and evaluation since it can take advantage of mini-batch optimization , and learns to adapt their priors to data . The inferred dependency graph in our model construction can be used to visualise the high dimensional latent representation , obviating the need for ad-hoc post-processing steps as required in most of the multi-modal learning methods . The presented work makes four major contributions : 1 ) We develop a novel multi-modal Relational Neural Process , mRNP , that defines a distribution over functions of multiple data types by employing local latent variables , and learns a dependency structure among the samples of the given modalities . 2 ) We theoretically prove exchangeability and consistency of mRNP , two necessary conditions that have to be satisfied during the construction of such a model , showing that mRNP is a valid stochastic process . 3 ) We further show that the local latent variable structure in mRNP is able to encode inductive biases and demonstrate this by designing an mRNP model that behaves similarly to a GP with an RBF kernel ( an ablation study ) . 4 ) We introduce mixture-of-graphs ( MoG ) in our model construction that can address the issues such as computational complexity and miscalibrated precisions observed in the previous approaches in multi-modal learning . 2 METHOD . We propose a new graph-structured multi-modal learning method , referred as multi-modal Relational Neural Process ( mRNP ) , that combines the benefits of neural networks with that of stochastic processes . In the supervised multi-modal setting , our dataset contains sets of feature-label pairs , { x̄ , y } from |V| different modalities with various missing patterns , where |V| is number of views , x̄ = { xv|vth modality present } v∈V ∈ X̄ is the input modalities and y ∈ Y is the given label . One of the key motivation of mRNP is that it gives us the ability to define a distribution over functions rather than learning a single static function when fitting the data . In this framework , inspired by the idea of inducing points in sparse Gaussian processes ( Titsias , 2009 ; Damianou & Lawrence , 2013 ) , we first select a reference set of samples that consists of functions f : x̄→ y that are sampled from some underlying distributionH . We then establish a probability distribution over fh around those samples . We define the reference set , XR = { XR,1 , . . . , XR , V } ⊂ X̄ , where XR , v ∈ RNv×Dv , and 1Rapid adaption to new observations is what distinguishes a model that can learn a whole distribution over functions by seeing different functions in each batch from the one that learns a single underlying function by observing several batches during the training . At the test time , the former will narrow down what the current function is by seeing a few context observations . Nv and Dv are the number of samples and dimension of features in the domain v , respectively . We also define O = X̄ \XR , which is the set of all possible samples that are not in XR . Please note that two sets XR and O includes both labeled and unlabeled samples . We also denote our labeled samples as Dx = { x̄1 , . . . , x̄N } , a set of any finite random set from X̄ , as well as defining sets XM = Dx\XR and XB = XR ∪ XM . The Venn diagram of the sets are provided in Figure 1 . We infer the parameters of our probabilistic model using variational inference . In the following , we first introduce each of the latent variables in our model as well as their corresponding prior and posterior distributions , and then prove that it corresponds to an infinitely exchangeable stochastic process . The graphical model of mRNP is illustrated in Figure 1 . Modality-specific latent spaces . The first step is to construct modality-specific latent representations , denoted by U = { Uv } v∈V , a set of Nv ×Du matrices , which independently embed samples of each modality ( i.e . Nv ) into a Du dimensional latent space as follows : ∫ pθ ( UB , XB ) dX = ∏ v∈V ∫ pθ ( UB , v , XB , v ) dXB , v = ∏ v∈V ∫ pθ ( UB , v |XB , v ) p ( XB , v ) dXB , v . We further parameterize the distribution over the samples ui , v independently : pθ ( UB , v |XB , v ) = ∏ i∈Bv pθ ( ui , v |xi , v ) , ( 1 ) where pθ ( ui , v |xi , v ) can be any distribution that is derived from the observed input data . For simplicity , we use a diagonal Gaussian , the parameters of which are a function of the observed input . More specifically , we use two functions denoted by ϕemb , µv ( XB , v ) and ϕ emb , σ v ( XB , v ) to infer the mean and variance of the distribution for each modality . Depending on the nature of the modality , these functions can be implemented using any highly expressive functions such as many variants of neural networks ( NNs ) and graph NNs . A directed graph among the reference samples . The next step is to construct a directed graph across the samples inXR using different modalities . Given the latent embedding UR , we first construct a set of directed acyclic graphs ( DAGs ) , G = { Gv } v∈V , where Gv is a random binary adjacency matrix between reference samples present in the vth modality . Then we combine the graphs Gv from multiple views to construct a common graph G. Inspired by the concept of stochastic orderings ( Shaked & Shanthikumar , 2007 ) , we impose a topological ordering over the vectors in UR to avoid cycles in each individual Gv . The distribution of the adjacency matrices are defined as follows : p ( Gv |UR , v ) = ∏ i∈Rv ∏ j∈Rv , j 6=i Bernoulli ( Gij , v | I [ t ( ui , v ) > t ( uj , v ) ] ϕsim ( ui , v , uj , v ) ) . ( 2 ) Similar to Louizos et al . ( 2019 ) , we employ a parameter free scalar projection t ( ui , v ) = ∑ k tk ( uik , v ) , where tk ( · ) is a monotonic function . In practice , we use the log cumulative distribution function of a standard normal distribution . Given the graphs Gv , we construct a directed graph G , in which the weight of each edge is average of its weights in different modalities , i.e . G = ∑ v∈V βvGv/ ∑ v∈Vmvm T v , where mv is a binary vector and mi , v = 1 when sample i is observed in the modality v , and βv = 1/|V | , assuming that the all modalities are equally important . The proposed formulation , referred as mixture-of-graphs ( MoG ) , is used to learn the shared structured representation of multi-modal inputs by using the graphs of individual modalities . A relational graph of dependencies from reference samples to XM . In addition to the directed graph G used for the reference samples , we apply MoG to construct a bipartite graph A = ∑ v∈V αvAv that is a bi-adjacency matrix from XR to XM . Given the latent embedding UB as described previously , we model the elements of each individual bi-adjacency matrix by Bernoulli distribution as follows : p ( Av |UM , v , UR , v ) = ∏ i∈Mv ∏ j∈Rv Bernoulli ( Aij |ϕsim ( ui , v , uj , v ) ) , ( 3 ) where ϕsim ( · , · ) is a score function measuring the similarity between the latent representations of the input samples in each modality v , and αv = 1/|V | , assuming the all modalities equally important . Depending on the desired relational inductive biases , we can appropriately define the function ϕsim ( · , · ) . The Bernoulli-Poisson link ϕsim ( ui , v , uj , v ) = 1− exp ( − ∑Du k=1 τk uik , v ujk , v ) and ϕsim ( ui , v , uj , v ) = exp ( − τ2 ||ui , v − uj , v|| 2 ) are two examples of such potential score functions . Please note that we use the latter one throughout this paper . For both A = { Av } v∈V and G , we use the concrete relaxation ( Maddison et al. , 2016 ; Gal et al. , 2017 ) during training while we sample from Bernoulli distributions in the testing phase . A shared latent variable . Having obtained the sample representations U and the dependency graphs A and G , we construct a shared latent variable , denoted by N ×Dz matrix Z , which can be used to predict target variable distributions yi . We parameterize the distributions of the predictive targets as follows : ∫ pθ ( yB , ZB |A , G , XR ) dZB = ∫ ∫ pθ ( yR , ZR |G , XR ) pθ ( yM , ZM |A , XR , yR ) dZR dZM = ∏ i∈R ∫ pθ ( zi |parGi ( XR , yR ) ) pθ ( yi | zi ) dzi ∏ j∈M ∫ pθ ( zj |parAj ( XR , yR ) ) pθ ( yj | zj ) dzj , where parGi ( · ) , parAj ( · ) are functions that return the parents of the points i and j according to G and A , respectively . We summarize the information from the parent samples and their targets in XR through the local latent variable zi , allowing the distribution yi to be explicitly dependent on the available graph dependency structures A and G. We set the distribution over zi as an independent Gaussian distribution whose parameters are a function of either A or G although any distribution with a permutation invariant probability density with respect to the parents can be used . More formally , the following distribution is adopted : pθ ( zi |parAi ( XR , yR ) ) = Dz∏ k=1 pθ ( zik |parAi ( XR , yR ) ) = Dz∏ k=1 N ( µpriorik , σ prior ik ) , ( 4 ) µpriorik = ci ∑ j∈XR Aijϕ prior , µ k ( xj , R , yj , R ) , σ prior ik = exp ( ci ∑ j∈XR Aijϕ prior , σ k ( xj , R , yj , R ) ) , where xj , R = { xj , R , v } v∈V , ϕprior , µ and ϕprior , σ are transform functions with a co-domain R|z| , and ci = ( ∑ j Aij + ) −1 is a normalization constant . Given input { XB , v } v∈V , we factorize the variational posterior distribution as qφ ( Z | XB ) = ∏ i∈B qφ ( zi |xi , B ) , ( 5 ) In practice , we also parameterize the priors over the latent zi in terms of the posterior distribution of the reference samples , i.e . qφ ( zi |xi , R ) . More precisely , we define ϕprior , µ and ϕprior , σ in equation ( 4 ) as ϕpost , µ ( xi , R ) +ϕ emb , µ label ( yi , R ) and ϕ post , σ ( xi , R ) +ϕ emb , σ label ( yi , R ) , respectively ; where ϕpost and ϕemblabel provide the means and variances of qφ ( zi |xi , R ) and the linear embedding of the labels of reference samples , respectively . We should point out that in this formulation , the prediction of yi indirectly depends on the input samples xi through the graphs G and A , which are functions of ui . This leads to an uninformative standard normal prior over zi for the samples with very small probability of being connected to the reference set through A , and hence the prediction of yi will be constant . This can be seen as encoding inductive bias similar to a GP with an RBF kernel . MoG vs PoE . PoE is a prime choice to learn the variational joint posteriors ; it can learn under any combination of missing modalities ( Shi et al. , 2019 ) . Despite this advantage , it also has two main limitations : 1 ) the underlying model suffers from the overconfident experts since experts with greater precision will have more influence over the combined prediction than experts with lower precision ; 2 ) training and inference can be costly due to its artificial sub-sampling procedure . Given the aforementioned limitations of PoE , we propose to build a graph of dependencies among local latent variables as a mixture of modality-specific graphs , i.e . mixture-of-graph ( MoG ) , and then parameterize the joint latent space Z conditioned on those graph . Unlike PoE , MoG effectively spreads its density over all the individual modalities by imposing a weight to each individual graphs . When prior knowledge is available , terms αv and βv in the MoG construction can be used to encourage certain modalities , which is the case in multi-omics data integration where individual modalities often contain complementary information of target task . A more detailed discussion can be found in Section D of the supplement . We also should point out that mRNP is not a graph learning model . Rather , we propose a novel multi-view NPs that is able to solve a pitfall in utilizing stochastic processes for multi-view setting by learning a graph of dependencies . We can also consider graph learning in mRNP as a kind of cross-attention in the form of a dependency graph among local latent variables . From this viewpoint , mRNP is learning attention mechanism to combine modalities . The current framework in which mRNP learns the individual graphs might not be computationally efficient , especially if training set is large . We may improve the scalability issue of graph construction part of our mRNP by e.g . considering hierarchical structure similar to hFGW ( Xu et al. , 2020 ) . We leave this for a future study . Overall likelihood and learning . Putting everything together , the marginal likelihood is pθ ( yB | XB ) = ∑ G , A ∫ ∏ v∈V pθ ( UB , v |XB , v ) p ( G , A | UB ) pθ ( yB , ZB |A , G , XR ) dUB dZB . ( 6 ) We should point out that , for the sake of simplicity , we choose the reference set to be part of the training set Dx throughout this paper . More specifically , we assume XB = Dx . In case of existing reference samples that are not part of the training set , we need to marginalize out over them in equation ( 6 ) in addition to the marginalizations over the latent variables and structured graphs . We can also use the marginalization technique to incorporate unlabelled data in order to learn a better representation and/or to impute the missing labels . We leave this for a future study . We deploy variational inference to optimize the model parameters θ and variational parameters φ by minimizing the following derived Evidence Lower Bound ( ELBO ) : L = Eqφ ( ZR | XR ) p ( G | UR ) ∏v∈V pθ ( UR , v |XR , v ) [ log pθ ( yR , ZR |G , XR ) − log qφ ( ZR | XR ) ] + Eqφ ( ZM | XM ) p ( A | UB ) ∏ v∈V pθ ( UB , v |XB , v ) [ log pθ ( yM |ZM ) + log pθ ( ZM |parA ( XR , yR ) ) − log qφ ( ZM | XM ) ] . Please note that we can use mini-batches for the second expectation term where the size of the batches scale according to the size of the reference set XR . Due to the DAG structure in the first expectation term , we can not decompose it to independent sums . Predictive distribution . After optimizing the model parameters , we derive the predictive distribution for unseen samples { x∗v } v∈V with missing modalities based on Bayes theorem as follows : pθ ( y ∗ | { x∗v } v∈V , XB , yB ) = ∑ a∗ ∫ ∏ v∈V pθ ( u ∗ v , UR , v |x∗v , XR , v ) p ( a∗ | { u∗v } v∈V , UR ) pθ ( z∗ |para∗ ( XR , yR ) ) pθ ( y∗ | z∗ ) dUR du∗v∈V dz∗ , where u∗v is the embedding representation of the observed modality x ∗ v through the neural network ϕembv and a ∗ is a binary vector similar to a row of A that denotes reference parents of the new samples . This is similar to the predictive distribution in few-shot learning ( Sung et al. , 2018 ) . Proposition I. mRNP corresponds to Bayesian models as the distributions defined in Equation ( 6 ) are valid , permutation invariant stochastic processes . proof . The proof is deferred to the supplement .

Summary: This paper tries to use Bayesian relational generative model for scalable multi-modal learning. They propose a class of stochastic processes that learns a graph of dependencies between samples across multi-modal data types through adopting priors over the relational structure of the given data modalities. The so-called mRNP method can address the limitations in joint posterior approximation.