Scalable Permutation Invariant Multi-Output Gaussian Processes for Cancer Drug Response

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-06 13:19:43

Abstract: Dose-response prediction in cancer is a critical step to assessing the efficacy of drug combinations on cancer cell-lines. The efficacy of a pair of drugs can be expressively modelled through a dose-response surface which outputs the viability score across a spectrum of drug concentrations for each pair of drugs in the training data. Using large in-vitro drug sensitivity screens, the goal is to develop accurate predictive models that can be used to inform treatment decisions by predicting the efficacy of given drug combination on new cancer cell lines as well as predict the effect of unseen drugs. Previous work, proposed a framework for modelling dose response surfaces with multi-output GPs, however, the model relied on the exact GP marginal likelihood and prohibited scalable inference. Further, the only inputs were drug concentrations per pair while the triplet of cell-lines and drug pair corresponded to different outputs . We make two important innovations in this work, we propose a framework for stochastic multi-output GPs for scalable inference; and, use a deep generative model (DGM) to embed the drugs in a continuous chemical space - enabling viability predictions for unseen drugs. We demonstrate the performance of our model in a simple setting using a high-throughput dataset and show that the model is able to efficiently borrow information across outputs.

Main paper: 
 introduction: 
 In recent years, much work has gone in to developing predictive models for cancer treatments. Using data from high-throughput in-vitro experiments on multiple drugs and cell lines, large models are trained with the goal of predicting the sensitivity of a drug on a certain cell line. In the context of drug combinations, interest has frequently been on predicting a summary measure of drug interaction, e.g. a synergy score [11] , computed from fitted dose-response surfaces, and an assumption of how non-interacting drugs should behave. Synergy scores are inherently quite crude measurements of drug interaction, and fundamentally hinge on the choice of non-interaction assumption. For this reason, several authors [7, 18, 15, 5] have proposed algorithms that instead aim to predict the entire dose-response surface. Since dose-response experiments are naturally invariant to the ordering of the drugs in a combination, we further augment the model by directly encoding this invariance into the prior.In this paper, we build on previous work on permutation invariant multi-output Gaussian Processes (PIMOGPs) by [15] ; the permutation invariance arises naturally in the context of dose-response functions as these are invariant to the ordering of the drugs in a combination. The original PIMOGPs framework relied on using exact GPs and thus incurred cubic complexity. Further, it did not encode the drug pairs i.e. each drug combination experiment (cell line, drug A, drug B) triplet was considered an output, and only the concentrations (c A , c B ) were given as inputs. In this work we instead take only the cell line as output, and regard the drugs as inputs alongside the drug concentrations. In order to encode the drug information as inputs we make use of a deep generative model that takes as input a string representation of the molecule, and outputs a low-dimensional representation of the drug (see A).The main contribution of this work is deriving the stochastic variational bound for the permutation invariant linear model of coregionalisation (LMC). This method naturally handles missing data and provides uncertainty quantification. 
linear model of coregionalisation for mogps: 
 Multi-output GPs (MOGPs) are the extension of GP regression to the setting where the regression outputs are multidimensional, i.e. for any input x, the resulting mapping f (x) ∈ R m for some m > 1. There are many ways of constructing MOGPs (see [2] for a review), but our focus here will be on the linear model of coregionalisation (LMC).In LMC the outputs are modelled as linear combinations of a set of independent latent functions, that are themselves modelled as GPs. That is, considering a set of m outputs {f j (x)} m j=1 for an inputx ∈ R p , the j-th output is modelled as,EQUATIONwhere u r ∼ GP(0, k r (•, •)) for r = 1, . . . , R independently, and a j1 , . . . , a jR are scalar weights.Note that each latent function is given its own covariance function k r , but these are free to have the same covariance, while maintaining independence. Latent functions that share their covariance functions can be grouped into G groups with R g latent functions in each group, and the equation rewritten as:f j (x) = G g=1 Rg r=1 a (r)jg u (r) g (x).Since GPs are closed under addition, this construction induces a GP over all outputs. The crosscovariance between two evaluations f j (x) and f j ′ (x ′ ) can be written asCov [f j (x), f j ′ (x ′ )] = G g=1 Rg r=1 a (r) jg a (r) j ′ g k g (x, x ′ ) = G g=1 b (g) jj ′ k g (x, x ′ ),where b(g) jj ′ = Rg r=1 a (r)jg a (r) j ′ g . The full cross-covariance over all m outputs can be written asEQUATIONwhere the matrix B g is known as a coregionalisation matrix, with entries{B g } ij = b (g) ij .In the case of a complete dataset where every input is observed at every output, the above expression can be written using Kronecker products:EQUATIONwhere K g has entries {K g } ij = k g (x i , x j ). In the simplest setting, where all latent functions share the same covariance function (i.e. G = 1), this model is known as the intrinsic coregionalisation model (ICM) and the Kronecker structure can be exploited to obtain large computational gains.For the LMC, the matrices are even larger due to the covariance across outputs, resulting in a cubic complexity in both inputs and outputs -O(n 3 m 3 ). In the special case of the IMC, i.e. where G = 1 in equation ( 3), using various Kronecker tricks can bring this down to O(n 3 + m 3 ) -see e.g. [16] .3 Stochastic Variational Linear model of Coregionalisation for Multi-outputIn this section we desribe our main insight of leveraging SVI in the setting of the LMC in order to speed up the necessary computations.Instead of having a single set of inducing inputs, we allow different inducing locations per latentfunction Z = {Z r } R r=1, where Z r = {z 1r , . . . , z qr }, with corresponding inducing variables for eachlatent function U = [u T 1 , . . . , u T R ] T, where u r = [u r (z 1r ), . . . , u r (z qr )] T . In order to not overload the notation, we assume the same number of inducing variables, q, per latent function, and that the dataset is complete, i.e. that every input x is observed at every output. We denote by Y and F vectors of observations and latent evaluations, stacked for each output, i.e. Y = [y T 1 , . . . , y m ] T where y j = [y 1j , . . . , y nj ] T , where y ij = f j (x i ) + ϵ ij , and similarly for F.The variational lower bound in the LMC case takes the same form as the stochastic variational bound for single output regression (see A.1), factorising over observations and outputs, but the matrices involved are more structured:EQUATIONwhereα T ij = A j K (i) fu K -1uu , A is the matrix with LMC coefficients, i.e. {A} ij = a ij , where A j denotes the j-th row of A. The matrixK (i)fu is an R × Rq block-diagonal matrix with entries {k r (x i , Z r )} R r=1 , i.e. formed by evaluating the input point x i against all the inducing points Z r in the r-th latent function. Qij is formed byA j [K (i) ff -K (i) fu K -1 uu K (i)uf ], where the matrixK (i) ff is an R × R block-diagonal matrix with entries {k r (x i , x i )} R r=1 , K (i) uf = (K (i) fu ) T and K uu is an Rq × Rq block-diagonal matrix with entries {k r (Z r , Z r )} R r=1 .Finally, μ and Σ come from the variational distribution q(U) -which we assume factorise over the r latent functions q(U) = R r=1 q r (u r )where q r (u r ) = N (µ r , Σ r ), yielding a block-diagonal Σ, and a decomposition of the KL-divergence term as KL(p(U)∥q(U)) = R r=1 KL(p r (u r )∥q r (u r )). 
permutation invariance: 
 For the problem of dose-response prediction, interest is on encoding an invariance on the ordering of the drugs -a permutation invariance of the inputs. In the context of multi-output GPs, we want to encode a permutation invariance for every output, f j . Looking at the construction of the LMC in equation ( 1), it suffices to ensure that each latent function u r has the required invariance. Lettingx denote a permuted version of the input x, with the desired invariance to encode u r (x) = u r ( x), this can be achieved by introducing another function ũr (x) and constructing u r (x) via a summation argument:EQUATIONfrom which we see that the mapping x → x leaves the function unchanged. Placing a zero-mean GP prior on ũr with kernel kr (•, •) induces a zero-mean GP prior on u r with kernel,EQUATIONIn order to enable SVI in the context of permutation invariant MOGPs, a slight modification needs to be made regarding the inducing variables u r . Specifically, instead of regarding u r as observations from the latent function u r , they are assumed as observations from the underlying ũr . The structure of the bound in equation (4) remains unchanged. Only, the entries of matricesEQUATIONfu and K uu need to be computed according to the following equations:EQUATIONK uu : k r,uu (Z r , Z r ) = kr (Z r , Z r ). We use the data from [12] , and follow the same pre-processing procedure as in [15] . This dataset We test the performance of our model in the leave-triplet-out (LTO) setting, using the nomenclature of 120 [1] . That is, we consider prediction of a specific (cell line, drug A, drug B)-triplet that does not appear We used a batch size of 256, and trained for 12 epochs using the Adam optimizer [8] with an initial 131 learning rate of 1e-2 decreasing to 1e-3 after 6 epochs, and to 1e-4 after 9 epochs. We plot the training 132 loss over the epochs in Figure 1 (Left). As performance metrics, we compute the root mean squared normalised it to a correlation matrix. We note that in contrast to the model of [15] , our implementation 141 is able to borrow strength across cell lines in its predictions. Simplifying the KL between the variational and true posterior -KL(q(f , u)||p(f , u|y)) yields the following lower bound.log p(y|θ) ≥ E q(f ,u) [log p(y|f )]-KL(q(u|θ)||p(u|θ)).If the first term above entails a factorisable likelihood, then the expectation w.r.t the variational distribution q(f , u) = p(f |u, θ)q(u|θ) is analytically tractable yielding,EQUATIONwhere k T n is the n th row of K fu and knn is the n th element on the diagonal of the matrix K ff -K fu K -1 uu K uf ; both these terms are only dependent on their respective data-point x n .In the PIICM, each drug combination experiment (cell line, drug A, drug B) triplet was considered an output, and only the concentrations (c A , c B ) was given as inputs. In this manuscript we instead take only the cell line as output, and regard the drugs as inputs alongside the drug concentrations. In order to encode the drug information as inputs, we could one-hot encode them, but instead make use of a deep generative model that takes as input a string representation of the molecule, and outputs a low-dimensional representation of the drug.Most of the literature on machine learning based chemical design for the string representation of molecules uses SMILES strings [19] -a line notation method which encodes molecular structure using short ASCII strings. However, the SMILES representation has two critical limitations.First, they are not designed to capture molecular similarity, hence molecules with almost identical structure can have markedly different SMILES [6] . Second, they are not robust on their own, which means that generative models are likely to produce strings that do not represent valid molecules.Hence, the latent space of DGMs trained on SMILES strings can potentially have large dead zones where none of the points sampled in the region decode to valid molecules. To overcome these issues, we train our model on an alternative string representation for molecules introduced The deep generative model used in this context is trained for autoencoding and features an openended chemical latent space learnt by embedding discrete molecules in a continuous vector space (encoder). For generation, an inverse step (decoder) converts a continuous vector in latent space to a valid molecule. This is the classical encoder-decoder set-up as in a standard VAE [9] . We use a recurrent VAE architecture with an RNN encoder and a decoder to sequentially process the 
 
 Appendix:  



Reviews: review: title: title: Sound approach to scalable MOGPs
main_review: main_review: This paper presents a novel and technically sound approach to scalable MOGPs for cancer drug response prediction, incorporating permutation invariance and deep generative models. While the results are promising, the complexity of the method and limited generalizability are difficult for me to understand.

Pros:
1. The paper highlights a new efficient method by using stochastic variation inference.
2. The results are well presented and the paper is well written.
3. The paper highlights an interesting use of deep generative models to embed drugs as part of a vector space.

Cons:
1. Difficult to understand the intricacies of the generative model.
2. Hard to understand how it will scale and what obstacles lie in practical use cases.
title: title: an interesting paper with detailed formulation of the model. while lacks some details in experiment results
main_review: main_review: The paper presents a technically sound and innovative approach to dose-response prediction in cancer drug combinations using permutation-invariant multi-output Gaussian Processes (MOGPs). It effectively leverages a combination of Stochastic Variational Inference (SVI) and Deep Generative Models (DGMs) to overcome key limitations in previous works. The methodology is mathematically rigorous, employing well-established Gaussian process frameworks like Linear Model of Coregionalization (LMC) while extending them to more scalable and flexible models.
The originality of this work is high, as it introduces new solutions to known problems in cancer drug-response prediction:
	•	Permutation Invariance: Encoding permutation invariance directly into the Gaussian process framework is a notable innovation that respects the inherent symmetry in dose-response experiments.
	•	Scalability: By leveraging Stochastic Variational Inference (SVI), the paper improves scalability, addressing the prohibitive cubic complexity of traditional GP models.
	•	Use of SELFIES: Incorporating the SELFIES molecular representation to improve the robustness of chemical embeddings in drug-response models is another innovative contribution.

The only part which is not satisfying is the relatively short result sections. It could benefit from more detailed analysis and, crucially, comparisons with other SOTA methods. What is the performance compared with others? If it is not as good as others, whct is the benefits and potential strength of the proposed method?