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 proposed to use neural processes for supervised multi-modal learning, which have the abilities to estimate the uncertainty of prediction and to handle missing modalities. Specifically, a directed acyclic graph is learned for each modality in the neural process, which are then used to construct a mixture-of-graphs (MoG) to sidestep the modalities missing problem. The introduction of inducing points makes the predictive distribution tractable. Experimentally, on the label prediction and uncertainty estimation tasks, the proposed model performs well compared to the recent multimodal learning methods.