Luwei Wang

School of Informatics, University of Edinburgh, Edinburgh, UK, N\{\mathbf{x}_{n}\}_{n=1}^{N} input covariates, it is a multivariate Gaussian on these finite number of outcomes 𝝁0∈ℝN{\boldsymbol{\mu}}^{0}\in\mathbb{R}^{N} conditioned on their covariates: 𝝁0|𝐗∼𝒩N​(𝐦0,𝐊0​(𝐗)){\boldsymbol{\mu}}^{0}\,|\,\mathbf{X}\sim\mathcal{N}_{N}(\mathbf{m}_{0},\mathbf{K}_{0}(\mathbf{X})) where 𝐦0=𝟎\mathbf{m}_{0}=\mathbf{0} is an NN-vector of unknown function values and 𝐊0​(𝐗)\mathbf{K}_{0}(\mathbf{X}) is an N×NN\times N covariance matrix. We use the squared exponential kernel with noise for the covariance function, with the marginal standard deviation (α\alpha), controlling the overall scale of function values and length-scale (ρ\rho), determining how far apart inputs must be before their function values become uncorrelated. To capture varying influence across covariates, we adopt separate length scales for each dimension, following the automatic relevance determination (ARD) approach [Neal, 1996]. Thus, GP parameters ϕ={α,𝝆}{\boldsymbol{\phi}}=\{\alpha,{\boldsymbol{\rho}}\}. Note that it is possible to use a linear kernel for GP to consider a simpler linear regression fit. Choices for Treatment Effect For better interpretability, we consider a constant treatment effect over each cluster τ​(𝐱;𝜷k)=τk=βk\tau(\mathbf{x};{\boldsymbol{\beta}}_{k})=\tau_{k}=\beta_{k}. This homogeneity assumption is analogous to conventional risk stratification approaches in clinical trials, where sample sizes are small and constant treatment effects are commonly assumed across risk strata [Kent et al., 2020]. We can also consider a linear treatment effect [Shahn and Madigan, 2017], i.e., τ​(𝐱;𝜷k)=𝜷k⊤​𝐱\tau(\mathbf{x};{\boldsymbol{\beta}}_{k})={\boldsymbol{\beta}}_{k}^{\top}\mathbf{x}, but focus on the former in this paper. Intuitively, we jointly cluster samples based on their covariates and their treatment effect as shown in Figure 1(b, right). The assumption of a normal distribution and a constant treatment effect facilitates interpretability and coherence by yielding clusters with spherical geometry and within each cluster, individuals share the same treatment-response characteristics. Feature Selection Following the approach by Liverani et al. [2015], we employ the cluster-specific soft feature selection approach. Each mixture component kk has an associated vector 𝜸k=(γk,1,…,γk,D){\boldsymbol{\gamma}}_{k}=(\gamma_{k,1},\dots,\gamma_{k,D}), where γk,d\gamma_{k,d} is a latent variable in (0,1)(0,1) that determines whether the dd-th covariate is important to mixture component kk. For the binary covariates, let θd​0p\theta^{p}_{d0} be the sample prevalence of binary covariate dd over the population, then we define the new composite cluster parameters as θ¯d​kp≔γk,d​θd​kp+(1−γk,d)​θd​0p.\bar{\theta}^{p}_{dk}\coloneq\gamma_{k,d}\theta^{p}_{dk}+(1-\gamma_{k,d})\theta^{p}_{d0}. This expression is substituted in place of θd​kp\theta^{p}_{dk} to provide the likelihood for the covariate model. We further consider a weakly informative prior for γk,d∼Beta​(0.5,0.5)\gamma_{k,d}\sim\mathrm{Beta}(0.5,0.5). Similarly, for continuous covariates, we define θd​0μ\theta^{\mu}_{d0} to be sample mean of continuous covariate dd, and define new cluster parameters as θ¯d​kμ≔γk,d​θd​kμ+(1−γk,d)​θd​0μ.\bar{\theta}^{\mu}_{dk}\coloneq\gamma_{k,d}\theta^{\mu}_{dk}+(1-\gamma_{k,d})\theta^{\mu}_{d0}. 𝝅{\boldsymbol{\pi}}