# On Learning Prediction-Focused Mixtures

Abhishek Sharma  SEAS, Harvard University  Catherine Zeng  SEAS, Harvard
University  Sanjana Narayanan  SEAS, Harvard University  Sonali Parbhoo  SEAS,
Harvard University  Finale Doshi-Velez  SEAS, Harvard University

###### Abstract

Probabilistic models help us encode latent structures that both model the data
and are ideally also useful for specific downstream tasks. Among these,
mixture models and their time-series counterparts, hidden Markov models,
identify discrete components in the data. In this work, we focus on a
constrained capacity setting, where we want to learn a model with relatively
few components (e.g. for interpretability purposes). To maintain prediction
performance, we introduce prediction-focused modeling for mixtures, which
automatically selects the dimensions relevant to the prediction task. Our
approach identifies relevant signal from the input, outperforms models that
are not prediction-focused, and is easy to optimize; we also characterize when
prediction-focused modeling can be expected to work.

##  1 Introduction

Generative models provide a powerful mechanism for learning data distributions
that can explain the data well and help predict specific outcomes of interest.
For example, for patients with sepsis we might be interested in grouping
clinical phenotypes and biological markers that are correlated with a clinical
outcome such as mortality [15]. In this context, a generative model elegantly
handles missingness —common in real-world datasets [4]— and posits a mechanism
for experts to understand where the model might be going wrong.

In this paper, we consider the task of learning a generative model that is
also useful for predicting outcomes on some downstream tasks. In particular,
we focus on settings of limited capacity, i.e. when the model does not have
enough parameters to fully model the input distribution. The limited capacity
setting is relevant when we want models that are small enough to inspect, or
in real-world settings where there will always be some level of
misspecification.

Several methods have been proposed to learn generative models with good
discriminative performance [7, 3, 8]. The key challenge is to balance the
trade-off between the generative modeling objective and the prediction
objective. The simplest two-step approach of maximizing the likelihood of the
observed data and using that model for prediction largely fails because the
learned model may not have modeled aspects of the inputs most useful for the
prediction task. Alternatives such as prediction-constrained models [9]
explicitly model the generative-discriminative trade-off in a joint
optimization, but are difficult to optimize in practice.

In our work, we manage the trade-off between explaining the data and the
predictive objective by explicitly using ‘switch’ variables that allow our
models to select only those dimensions that are relevant to the downstream
task, and enable accurate predictions even in settings where number of
components is misspecified. Unlike [9], the entire problem can be written as a
graphical model with the discriminative-generative trade-off being managed by
a single hyperparameter corresponding to the probability that a dimension is
relevant. This setting was partially explored in [14], who considered a
specific context of relevant words for topic models; we add the following
contributions:

We extend the idea of managing the generative-discriminative trade-off via
relevant dimensions to continuous settings (which do not have the
orthogonality conditions used in [14]): Gaussian Mixture Models (GMM) and
Hidden Markov Models (HMM). We demonstrate our approach not only yields
performance benefits, but it is also more robust to optimize than more generic
prediction constrained approaches.

We provide evidence that our model achieves the desired generative-
discriminative trade-off directly from the maximum likelihood parameters of
the associated graphical model, instead of appealing to lower bounds or
inference assumptions (as in [14]). In doing so, we also illuminate when this
relatively simple, robust approach will work, and in what regimes it will
fail.

##  2 Related Work

Semi-supervised Learning aims to augment generative modeling of inputs X with
a supervisory signal $Y$ usually by training the joint objective $\log
p(\textbf{X},Y)$ [13, 11, 8]. When modeling the generative process with latent
variables Z as $\textbf{X}\leftarrow\textbf{Z}\rightarrow Y$, this approach
fails to recognize the inherent asymmetry between the inputs and the targets,
and therefore ignores the target dimension as the model allocates its capacity
to the more structured inputs. We address this issue by incorporating specific
input-focused latent variables that treat the inputs and targets differently.
Methods that incorporate an external signal to aid clustering of data exist
under the paradigm of “supervised clustering” but they fail to provide a
density of the observed data.

Task-focused Generative Modeling Many recent works have considered ways to
learn generative models that are useful for some downstream tasks [12, 5, 16,
6]. Most of these train for the discriminative loss only. Closer to our goal
of managing the generative-discriminative trade-off, [9] proposed prediction-
constrained learning in the context of mixture and topic models, noting the
asymmetry of the task in predicting labels from documents (and not the other
way around). Their framework balances the maximum likelihood objective with
predictive performance by introducing the latter as a constraint and then
solving the resulting constrained optimization problem. But the resulting
constrained objective is often hard to optimize in practice, and unlike our
work, it does not correspond to maximum likelihood of a valid graphical model.

In the context of topic models for text documents, [14] propose using switch
variables to exclude or include words depending on their relevance for
prediction. Their approach relies on independence assumptions between switches
and topic distributions in the approximate posterior for the model to learn
useful topics. Our models do not rely on such conditions during inference, and
have the prediction-focused properties by virtue of their structure alone.

Information Bottleneck (IB) methods take a joint distribution
$p(\textbf{X},Y)$ and compress the inputs X into a latent code Z that is most
informative about the target $Y$ [17]. Both the traditional IB and the deep
variational IB formulations assume specific Markov structures for the
generative model [1, 19]. In contrast, we achieve the trade-off between
compressing inputs and predicting targets by proposing a specific
probabilistic model which encodes our knowledge about the generative process.
We demonstrate that this model is sufficient for getting the trade-off we
seek, independent of the IB objective.

##  3 Problem Setting

We study the case where the inputs X are a concatenation of $M$ independent
mixtures, each with $K_{m}$ components ($m\in[M]$). Without loss of
generality, we assume that the first mixture also generates our target of
interest, $Y$:

$\displaystyle\textbf{X}=[\textbf{X}_{1};\dots;\textbf{X}_{M}]$
$\displaystyle[\textbf{X}_{1};Y]\sim\textsc{GMM}(K_{1},\Theta_{1});\quad\textbf{X}_{m}\sim\textsc{GMM}(K_{m},\Theta_{m})$

This is a setting which is pervasive in real world datasets, where there is
often much more structure in the input data X than is needed to make
predictions for a specific task (e.g. predicting $Y$). To perfectly model the
input distribution $p(\textbf{X})$, we would need a very large generative
model (with $K=\prod_{m}K_{m}$ components). Such a large model might be hard
to fit with limited data, and may be unnecessary if we are most interested in
capturing structure in the data that is relevant to the task.

For example, in a healthcare setting, different dimensions may correspond to
sets of measurements related to different conditions such as kidney disease,
mental health, or a broken leg. If we are most interested in predicting
outcomes related to recovery from a broken limb, then modeling variation in
broken leg dimensions might be essential, whereas modeling variation in kidney
disease may be less important (even if a chronic disease generates more data).

Below, we consider misspecification in the number of mixture components and
ensure that the model prioritizes the use of components to learn those that
are relevant. Formally, our setting pertains to the case where we are provided
with a fixed budget K ( $<\prod_{m}K_{m}$) of components we can use, and we
are required to (a) perform well on the task of predicting $Y$ at test time,
and (b) return a generative model of data that achieves as high a value of
$p(\textbf{X})$ as possible given that we also want to perform well on
predicting $Y$.

##  4 Prediction-Focused Mixtures

$Z^{(n)}$${\theta}$$X_{d}^{(n)}$$B,\pi$$Y^{(n)}$$\eta$$\xi_{d}$$p$$D$$N$ (a)

$Z_{1}$$Z_{2}$$\dots$$Z_{t}$$\dots$$Z_{T}$$A$${\theta}$$\xi_{d}$$X_{t,d}$$B,\pi$$Y_{t}$$\eta$$p$$D$
(b)

Figure 1: The (a) pf-GMM and (b) pf-HMM graphical models

We now present prediction-focused modeling, first by introducing prediction-
focused Gaussian Mixture Models (pf-GMM) (Fig. 1(a)) and then extending it to
time series setting by introducing prediction-focused Hidden Mixture Models
(pf-HMM) (Fig. 1(b)).

The generative process of the model is:

for each dimension $d$ do

Draw switch $\xi_{d}\sim\textbf{Bern}(p)$

end for

for each datum $n$ do

Draw latent state $Z^{(n)}\sim\textbf{Cat}(\theta)$

for each dimension $d$ do

if $\xi_{d}=1$ then

Draw ‘signal’ obs.$X_{d}^{(n)}\sim F_{X}(X_{d}^{(n)}|Z^{(n)};B_{Z^{(n)},d})$

else if $\xi_{d}=0$ then

Draw ‘noise’ obs. $X_{d}^{(n)}\sim F_{X}(X_{d}^{(n)};\pi_{d})$

end if

end for

Draw target $Y_{d}^{(n)}\sim F_{Y}(Y_{d}^{(n)};\eta_{Z^{(n)}})$

end for

where $\textbf{Bern}(p)$ is a Bernoulli distribution with parameter $p$,
$\textbf{Cat}(\theta)$ is Categorical distribution with parameter vector
$\theta$, and $F_{X},F_{Y}$ can be any distribution over the support of
$X_{d}^{(n)}$. In our paper, we choose the inputs to be Gaussian for
simplicity of exposition, and hence refer to the model as pf-GMM. pf-GMM
filters out ‘irrelevant’ dimensions by the use of switch variables
$\\{\xi_{d}\\}$, allowing it to use only a subset of dimensions to model along
with the targets. This property allows it to achieve good downstream
performance even when it is (inevitably) misspecified w.r.t. the data
generating process. Here, we assume that the input data dimensions can be
partitioned as being either ‘relevant’ or ‘irrelevant’, where ‘relevant’
dimensions are predictive of the target, $Y$.

##  5 Inference

Since pf-GMM is a valid graphical model, we can use the maximum likelihood
objective to do learning and inference. The log likelihood of the pf-GMM is:

$\displaystyle\log p(\textbf{X},\textbf{Y})$
$\displaystyle=\log\sum_{\textbf{Z},\boldsymbol{\xi}}p_{\theta}(\textbf{Z})p_{p}(\boldsymbol{\xi})p_{B,\pi}(\textbf{X}|\textbf{Z},\boldsymbol{\xi})p_{\eta}(\textbf{Y}|\textbf{Z})$
(1)

Computing Eq. 1 is generally intractable, but in Sec. 6, we demonstrate that
in lower dimensional settings (where the computation is tractable), maximizing
this likelihood gives good discriminative performance—suggesting that the pf-
GMM model itself balances the generative-discriminative trade-off.

In practice, we optimize a lower bound (the ELBO) to Eq. 1. We choose the
following form for the approximate posterior:

$\displaystyle q(\textbf{Z},\boldsymbol{\xi}|\Theta,\varphi)$
$\displaystyle=\prod_{n=1}^{N}q(Z^{(n)}|\textbf{X}^{(n)},Y^{(n)},\varphi,\Theta)\prod_{d=1}^{D}q(\xi_{d}^{(n)}|\varphi_{d})$

where $\Theta=\\{{\theta},B,\pi,\eta\\}$ refers to the model parameters. In
particular, note that we have added an additional index to the switch
variable; now $\xi_{d}^{(n)}$ is a function of the datum $n$ as well as the
dimension $d$. This will allow for a cleaner factorization; we model the
posterior Bernoulli distributions $q(\xi_{d}^{(n)}|\varphi_{d})$ with a tied
parameter $\varphi_{d}$ for dimension $d$ to still get a per-dimension
relevance.

The distribution $q(Z^{(n)}|\textbf{X}^{(n)},Y^{(n)},\varphi,\Theta)$ is
simply the posterior distribution conditioned on the variational parameters
$\varphi$ and can be computed as:

$\displaystyle q(Z=k|\textbf{X},Y,\varphi)$ $\displaystyle\propto
p_{{\theta}}(Z=k)\cdot p(\textbf{X}|Z=k,\varphi)\cdot p_{\eta}(Y|Z=k)$

At prediction time, the distribution $q(Z=k|\textbf{X},Y,\varphi)$ is computed
without the $p_{\eta}(Y|Z=k)$ term and the feature selection parameter
$\varphi$ helps select the right features in the term
$p(\textbf{X}|Z=k,\varphi)$.

This results in the following ELBO objective (full derivation in Supplement
Sec. 1):

$\displaystyle\text{ELBO}(q)$
$\displaystyle=\operatorname*{\mathbb{E}}_{q(\textbf{Z},\boldsymbol{\xi})}\left[\log
p_{\Theta}(\textbf{X},\textbf{Y}|\textbf{Z},\boldsymbol{\xi})\right]$
$\displaystyle\quad\quad\quad-\textbf{KL}\left(q(\textbf{Z},\boldsymbol{\xi}|\Theta,\varphi)||p_{\Theta}(\textbf{Z},\boldsymbol{\xi})\right)$
(2)

which is tractable to compute because the posterior factorizes over
$\boldsymbol{\xi}$ and Z. After this setup, the model can be trained using the
variational EM algorithm which alternates between learning the parameters
$\Theta$ and inferring the approximate posterior
$q(\textbf{Z},\boldsymbol{\xi}|\Theta,\varphi)$.

Similarly, we train pf-HMM using variational inference by using the following
approximate posterior, $q(\textbf{Z},\boldsymbol{\xi}|\Theta,\varphi)$:

$\displaystyle=\prod_{n=1}^{N}\left(\prod_{d}\prod_{t}q({\boldsymbol{\xi}}^{(n)}_{t,d}|\varphi_{d})\right)p({\textbf{Z}}^{(n)}|{\textbf{X}}^{(n)},{\textbf{Y}}^{(n)},\varphi,\Theta)$

where the subscript $t$ indexes time within a sequence and
${\textbf{Z}}^{(n)},{\textbf{X}}^{(n)},{\textbf{Y}}^{(n)}$ are variables
corresponding to sequence $n$. Training proceeds using variational EM, with
the only exception that we now use the Forward-Backward algorithm to infer
$p({\textbf{Z}}^{(n)}|{\textbf{X}}^{(n)},{\textbf{Y}}^{(n)},\varphi,\Theta)$
during the E-step. We provide the complete derivation for the ELBO in
Supplement Sec. 1.

##  6 Analysis

(a)

(b)

(c)

(d)

(e)

(f)

Figure 2: Comparison of pf-GMM with a supervised GMM (sup-GMM) on a synthetic
dataset. (a) First 2 (out of 6) dimensions of the input data, with colors
identifying the target label. (b) Learned components of both models when the
component budget is 4 (identical in this case). (c) Learned components of both
models restricted to 2 components—pf-GMM (red contours) learns components that
are predictive of the labels whereas sup-GMM (blue) components achieve higher
$p(\textbf{X})$ because they align with more input dimensions. (d) Heldout $Y$
AUROC (predictive performance) vs component budget (K)—pf-GMM performs better
under misspecification and no worse than sup-GMM otherwise. (e) Heldout AUROC
vs pf-GMM hyperparameter $p$—performance peaks at an optimal $p$ and drops
when it is higher (more dimensions considered as signal) or lower (no
dimension considered as signal). (f) Inferred switch posterior for the case
when component budget is 2 (4-component case is the same)—pf-GMM identifies
the relevant dimension.

Earlier work motivated using a switch-variable based approach to relevant and
irrelevant dimensions by noting that Eq. 1 can also be lower bounded by the
following [14] (Supplement 9.4):

$\displaystyle\log p(\textbf{X},\textbf{Y})$
$\displaystyle\geq\operatorname*{\mathbb{E}}_{p(\boldsymbol{\xi})}\left[\log
p(\textbf{Y}|\textbf{X},\boldsymbol{\xi})\right]+p\operatorname*{\mathbb{E}}_{p(\textbf{Z})}\left[\log
p_{B}(\textbf{X}|\textbf{Z})\right]$ $\displaystyle\quad+(1-p)\log
p_{\pi}(\textbf{X})$ (3)

While this bound does explicitly demonstrate a generative-discriminative
trade-off analogous to [9], the bound can neither be optimized directly in a
computationally efficient manner, nor does it imply that the joint likelihood
in Eq. 1 captures this trade-off—indeed, it would be unfortunate if our
results were because of our choice of approximate inference rather than the
due to the model itself.

In this section, we demonstrate using a synthetic example that the maximum
likelihood estimate (MLE) of the actual pf-GMM performs at least as well as
the GMM at the prediction task. We also characterize a limitation of
prediction-focused modeling that provides insight into the data regimes in
which we would expect this style of modeling to work. Specifically we show
that (a) the pf-GMM solution is equivalent to the GMM solution when there is
no misspecification, (b) the pf-GMM MLE is better than the GMM MLE solution
when there is misspecification in the number of clusters, and (c) pf-GMM MLE
is no better than than the GMM MLE when noise dimensions cannot be modeled
using a simple distribution, $F_{X}(\cdot;\pi_{d})$.

We choose an example that is instructive of the setting where pf-GMM would do
well: (a) the input data has structured noise coming from a mixture of
Gaussians, and (b) only a small proportion of the data dimensions is ‘signal’.
This simulates a setting in which ‘noise’ competes with the ‘signal’ for
parameters, and naïve joint modeling of inputs and targets is not expected to
do well. We assume the data comes from the following generative process:

$\displaystyle[X_{1},y]\sim
0.5[\mathcal{N}(0,1),\textbf{Bern}(0)]+0.5[\mathcal{N}(0,1),\textbf{Bern}(1)]$
$\displaystyle\textbf{X}_{2}\sim
0.5\mathcal{N}(\textbf{0},I)+0.5\mathcal{N}(\textbf{6},I);\quad\textbf{X}_{2}\in{\mathbb{R}}^{4}$
$\displaystyle\textbf{X}=[\textbf{X}_{1},\textbf{X}_{2}]$

where 0 and 6 are constant vectors in ${\mathbb{R}}^{4}$. In this data, only
the first dimension of $X$ is correlated with output $y$ while the rest of the
dimensions are not. Fig. 2(a) plots the first two dimensions of independent
samples from this distribution. We clearly see that we need four clusters to
model the data well and also predict targets correctly. Next, we see the
nature of solutions that are achieved using pf-GMM and GMM for different
component budgets. Finally, we will characterize a data setting where we would
not expect pf-GMM model to do better than a GMM.

##### No misspecification

Clearly, 4 components are sufficient to model the inputs accurately in this
problem. Therefore, given a budget of 4 components, the GMM will identify each
of the emission distributions, and at test time, will predict the output
labels correctly. Also note that the pf-GMM with $p=1$ will be equivalent to
the GMM, which is what we see when we look at the learned components in Fig.
2(b).

##### Misspecification

If we are restricted to a budget of only 2 components, the GMM identifies
clusters that align with the majority of input dimensions rather than aligning
itself to the first dimension (Fig. 2(c)). This is true because it is trying
to do well on the joint likelihood,
$p(\textbf{X},y)=p(\textbf{X})p(Y|\textbf{X})$. In the case of GMM, the
benefit from modeling $p(\textbf{X})$ well is much higher than cost of
modeling $p(Y|\textbf{X})$ poorly. The pf-GMM bypasses this issue by tuning
$p$ to a value that selects only a subset of dimensions to be relevant and
explains the rest using its noise distribution, $F_{X}(X_{d};\pi_{d})$. We can
therefore see in Fig. 2(e), that for an appropriate value of $p$, it selects
the clusters that capture the variation along outputs, leading to high
predictive accuracy.

##### Limitation of pf-GMM

The hyperparameter $p$ is useful in tuning the number of effective dimensions
that pf-GMM chooses to model using its components, trading off generative
performance with discriminative performance. Nonetheless, this causes pf-GMMs
to have a limitation which we characterize next. Suppose the irrelevant
dimensions were instead generated using the following distribution:

$\displaystyle\textbf{X}_{2}$ $\displaystyle\sim
0.5\mathcal{N}(\textbf{0},I)+0.5\mathcal{N}(\bm{\mu},I);\quad\textbf{X}_{2}\in\mathbb{R}^{4}$

where 0 and $\bm{\mu}$ are constant vectors in ${\mathbb{R}}^{4}$, and $\mu$
is the distance between the component means of each of the irrelevant
dimensions (previously, $\bm{\mu}$ was 6). Now consider the case when $\mu$
becomes large, i.e. the noise components provide a much higher contribution to
$p(\textbf{X})$ for any dimension. Let the parameters learned with and without
prediction-focus in the above example be $\Theta_{pf}$ and $\Theta$
respectively. Then as $\mu\rightarrow\infty$, their likelihood ratio goes to
0, making it asymptotically beneficial to select $\Theta$ as the maximum
likelihood solution. We prove this claim in Supplement Sec. 2, but the
intuition is that the pf-GMM selects the number of relevant dimensions.
However, which dimensions those are will depend on how much they contribute to
the generative term, $p(\textbf{X})$. If any dimension takes a much bigger hit
when modeled as noise (vs as signal) than other dimensions, pf-GMM will find
it hard to classify it as noise.

This analysis highlights in which ways a pf-GMM can be misspecified and still
give good predictive performance, and in which ways it cannot; specifically,
when each dimension can be reasonably modeled using a fixed Gaussian (i.e.
$\mathcal{N}(X_{d};\mathbb{E}[X_{d}],\mathbb{V}[X_{d}])$ does not go to zero),
pf-GMM will do well.

We note that the limitation of this model under this specific misspecification
is mitigated by the inference objective in Eq. 3, which also seeks out
parameters useful for predicting the targets for a larger range of $p$. We
show evidence for this in Supplement Sec. 3, but restrict our discussion and
experiments to the maximum likelihood objective here. Exploring this objective
is potentially interesting future work.

##  7 Experimental Details

(a) Synthetic

(b) HIV

(c) HAR

(d) Synthetic

(e) HIV

(f) HAR

Figure 3: Target AUROC performance as we vary the component budget $K$ in each
of the domains for GMMs (Top Row) and HMMs (Bottom Row) respectively. pf-GMMs
and pf-HMMs consistently outperform methods that do not perform the
prediction-generation trade-off for different choices of $K$. The synthetic
input data for the GMM case has 20 (out of 100 total) signal dimensions, with
both signal and noise dimensions created using independent GMMs of 10
components each. For the HMM case, the corresponding values are 2 (out of 20)
and 4 respectively.

In this section, we demonstrate that the pf-GMM can select relevant signal
from the input in both misspecified and fully specified settings. We
demonstrate that the model structure—involving the switch parameters and their
prior—is able to recover useful dimensions and maintain downstream performance
even under misspecification.

##### Baselines.

We compare the generative and predictive performance of pf-GMM to (1) an
approach that first learns a generative model using only the inputs and then
trains a logistic regression classifier on the posterior beliefs
$p(Z|\textbf{X})$ (2-Step-GMM), (2) a supervised Gaussian mixture model
without the switch parameters (sup-GMM), (3) a logistic regression model,
which is a discriminative approach that maximizes $\log p(Y|\textbf{X})$
directly (LogReg), and (4) a prediction-constrained GMM [9] (pc-GMM). For the
pf-HMM, we similarly compare to sup-HMM, 2-Step-HMM and LogReg.

##### Evaluation.

To evaluate prediction quality, we compare the AUROC for the categorical
target dimension on heldout data. To further assess models on their generative
and discriminative quality, we compute $\log p(\textbf{X})$ and $\log
p(\textbf{Y}|\textbf{X})$ on heldout data (computed without any approximations
to ensure consistent evaluation between the methods).

##### Datasets.

Synthetic. We use synthetic experiments to understand the behavior of mixture
models of varying degrees of misspecification. For our first experiment, we
test whether pf-GMM is robust to increasing the number of noise dimensions. To
simulate this setting, we vary the number of dimensions that come from the
‘signal’ mixture while keeping the number of dimensions fixed at 100. Next, we
test whether our model identifies the dimensions of interest under different
budgets of components. Finally, we test whether the inferred posterior values
for the switches indeed correspond to the true relevance of the dimensions.

We generate the signal and noise dimensions from independent GMMs, each with
10 components (we do the same for HMMs with 4 components). The emission
distribution of each mixture given the cluster identity K is
$p(\textbf{X}|Z=K)=\mathcal{N}(X;6K,I)$. The target is assigned a label in
$\\{0,1\\}$, with the assignment only depending on the cluster identities in
the signal mixture. The setup is presented in detail in Supplement Sec. 4.

(a)

(b)

(c)

Figure 4: For the above set of plots, signal and noise dimensions (100 in
total) were created using 2 independent GMMs of 10 components each; our
component modeling budget is 2. (a) Target AUROC as we vary the number of
signal dimensions in the data. When noise dimensions outnumber signal, pf-GMM
consistently outperforms methods which do not effectively trade off prediction
and generation quality. (b) Inferred switch posterior distribution when first
10 dimensions are relevant in the ground truth; pf-GMM correctly learns to
distinguish between signal and noise. (c) Variety in performance of optimal
hyperparameters of pc-GMM and pf-HMM under different component budgets (20
signal dimensions in data).

HIV. HIV is a virus that attacks the immune system which, if untreated, causes
AIDS. Therapy for HIV involves administering cocktails of antiretrovirals from
five classes namely, nnRTIs, nRTIs, PIs, FIs, and IIs to bring the viral load
below detection limits ($\leq 40$ copies/ml). Even in the presence of
treatment, the virus may mutate. Some mutations may cause drug resistance and
thus increase the viral load. These mutations are known as resistance-relevant
mutations and can be identified from a list of mutations by the IAS USA [18].
We study 53 236 patients with HIV from the EuResist Integrated Database [20].
Each person has a timeseries of approximately 16 steps associated with them.
Our task is to predict whether a treatment will bring the viral load below
detection limits in the next time-step. Each input contains 138 features
including CD4+ counts, mutations, treatments in terms of drug classes and lab
results. We compare the performance of pf-GMM and pf-HMM with their
corresponding baselines. Further details about the experimental setup are
provided in the Supplement Sec. 4.

Human Activity Recognition Dataset. The Human Activity Recognition [2] dataset
uses a smartphone to capture triaxial acceleration, angular velocity, and
other measurements of 30 participants while they perform one of 6 activities
of daily living. Each person has one time series associated with them, with an
average length of 338. Our task is to predict the daily activity the
participant performs at each timestep. We compare the predictive performance
of both pf-HMM and pf-GMM at different component budgets.

##  8 Results

##### Prediction-focused models identify parameters which perform better at
the downstream task.

pf-GMM and pf-HMM outperform their vanilla counterparts on target AUROC on
several real and synthetic datasets (Figures 2(d) and 3). This is because
their objectives maximize the data likelihood while paying attention to the
targets at training time. While the supervised variants also maximize the
joint likelihood of inputs and target, they treat the target as just another
input dimension. In contrast, pf– learning identifies the asymmetry in the
problem and puts the focus on input dimensions that are predictive of the
targets.

##### Prediction-focused models find relevant dimensions in misspecified and
high-noise settings.

Fig. 4(a) shows that the pf-GMM is able to make the best of its budget of two
components even when it is severely constrained (both signal and noise GMMs
had 10 components) and misspecified w.r.t the noise. This is possible because
the model learns to focus on the dimensions that matter, as can be seen from
its inferred switch posterior in Fig. 4(b). For HIV, we observe that for
patients where the pf– models predict that a treatment fails to bring the
viral load below detection limits, a number of resistance-relevant mutations
associated with that treatment are identified. As an example, for a patient
taking nRTIs experiencing virologic failure, pf– models identify certain
mutations such as K65R and M184V as relevant. These are consistent with [18].
Such insights may ultimately be useful for assessing treatment options as well
as offering alternative options for intervention.

##### Prediction-focused models balance the predictive and generative
objective.

To demonstrate that pf– models effectively trade off predictive and generative
objectives, we plot its performance on the discriminative vs generative
objective landscape (Fig. 5). Under severe misspecification, the sup-GMM and
2-Step-GMM models are not able to separate the signal from noise and thus pay
a price in discriminative performance. We also plot the pf-GMM model which is
directly trained on the discriminative objective, $\log p(Y|\textbf{X})$, and
see that focusing just on the discriminative term results in having poor
generative quality. pc-GMM also does well with respect to the trade-off,
because it explicitly incorporates $\log p(Y|\textbf{X})$ and $\log
p(\textbf{X})$ in its objective.

Figure 5: Discriminative-generative landscape. Here, 1 (out of 8) dimension is
relevant; both signal and noise were generated from different GMMs with 2
components; our modeling budget is 2 components

##### The switch posteriors can be interpreted as usefulness of the dimension
towards prediction.

Figures 2(f) and 4(b) demonstrate that the switch posterior identifies
dimensions that were relevant for prediction. Importantly, this property is
always true — not just when pf– models outperform their non-pf– counterparts.
None of the other methods assign an interpretable notion of relevance to the
inputs.

##### The switch prior is the only hyperparameter which we tune.

The only hyperparameter, $p$, is intuitive to tune because increasing it
monotonically increases our belief about how many relevant dimensions are
present. This also means that setting $p=0$ skips using the inputs altogether
and setting $p=1$ treats all dimensions as relevant. The optimal $p$ therefore
ends up somewhere in between, as also noted in Fig. 2(e).

##### Prediction-focused models are easier to optimize and achieve more stable
solutions.

pf– modeling benefits from being a valid graphical model and thus allowing
exact variational EM updates for the parameters. This means that our objective
is non-decreasing after every step—which is true for the vanilla methods but
not for pc-GMM. Because it has just one hyperparameter, searching over it is
just a 1-D search. While pc-GMM has just one hyperparameter $\lambda$—used for
trading off generative with predictive performance—it is still considerably
harder to optimize. The optimization requires a grid search over both learning
rates and the $\lambda$ hyperparameter, and takes much longer to converge than
the EM-based methods. We see from Fig. 4(c), that even after conditioning on
the best hyperparameters for each setting, pc-GMM solutions are much less
stable and would typically require more independent trials when compared to
pf-GMM solutions.

##  9 Discussion and Conclusion

In this paper, we study mixture modeling under misspecification and propose a
graphical model solution to the problem. This prediction-focused modeling
selects the ‘right’ features to cluster even when the inputs are corrupted
with several disparate noise mixtures. It achieves this trade-off by
maximizing the likelihood of the data directly, instead of resorting to any
property of a modified objective as proposed by [10] or of the approximate
posterior as done by [14]. We characterize the conditions where this model
outperforms its counterparts as well as the conditions where it does not. We
also show that prediction-focused models perform at least as well as
supervised mixture models and mixture models trained without any target, and
therefore are a useful alternative to these widely used models.

## References

  * [1] Alexander A Alemi, Ian Fischer, Joshua V Dillon and Kevin Murphy  “Deep variational information bottleneck”  In arXiv preprint arXiv:1612.00410, 2016
  * [2] D. Anguita et al.  “A Public Domain Dataset for Human Activity Recognition using Smartphones”  In ESANN, 2013
  * [3] David M. Blei and Jon D. McAuliffe  “Supervised Topic Models”  In Proceedings of the 20th International Conference on Neural Information Processing Systems, 2007  URL: https://arxiv.org/pdf/1003.0783.pdf
  * [4] Krzysztof J Cios and G William Moore  “Uniqueness of medical data mining”  In Artificial intelligence in medicine 26.1-2 Elsevier, 2002, pp. 1–24
  * [5] Adam D. Cobb, Stephen J. Roberts and Yarin Gal  “Loss-Calibrated Approximate Inference in Bayesian Neural Networks” arXiv: 1805.03901  In arXiv:1805.03901 [cs, stat], 2018  URL: http://arxiv.org/abs/1805.03901
  * [6] Joseph Futoma, Michael C. Hughes and Finale Doshi-Velez  “POPCORN: Partially Observed Prediction Constrained Reinforcement Learning”  In Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, 2020  URL: https://arxiv.org/pdf/2001.04032.pdf
  * [7] Kuzman Ganchev, Ben Taskar and João Gama  “Expectation Maximization and Posterior Constraints”  In Advances in Neural Information Processing Systems 20 Curran Associates, Inc., 2008  URL: https://proceedings.neurips.cc/paper/2007/hash/73e5080f0f3804cb9cf470a8ce895dac-Abstract.html
  * [8] Zoubin Ghahramani and Michael I Jordan  “Supervised learning from incomplete data via an EM approach”  In Advances in neural information processing systems, 1994, pp. 120–127
  * [9] Michael Hughes et al.  “Prediction-Constrained Training for Semi-Supervised Mixture and Topic Models” arXiv: 1707.07341  In arXiv:1707.07341 [cs, stat], 2017  URL: http://arxiv.org/abs/1707.07341
  * [10] Michael Hughes et al.  “Semi-Supervised Prediction-Constrained Topic Models”  In Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics 84, Proceedings of Machine Learning Research  PMLR, 2018, pp. 1067–1076
  * [11] Diederik P Kingma, Danilo J Rezende, Shakir Mohamed and Max Welling  “Semi-supervised learning with deep generative models”  In arXiv preprint arXiv:1406.5298, 2014
  * [12] Simon Lacoste–Julien, Ferenc Huszár and Zoubin Ghahramani  “Approximate inference for the loss-calibrated Bayesian”  In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics JMLR WorkshopConference Proceedings, 2011, pp. 416–424  URL: http://proceedings.mlr.press/v15/lacoste_julien11a.html
  * [13] Kamal Nigam, Andrew McCallum, Sebastian Thrun and Tom M. Mitchell  “Learning to Classify Text from Labeled and Unlabeled Documents.”  In AAAI/IAAI AAAI Press / The MIT Press, 1998, pp. 792–799
  * [14] Jason Ren, Russell Kunes and Finale Doshi-Velez  “Prediction Focused Topic Models via Feature Selection”  In International Conference on Artificial Intelligence and Statistics, 2020, pp. 4420–4429  PMLR
  * [15] Christopher W. Seymour et al.  “Derivation, Validation, and Potential Treatment Implications of Novel Clinical Phenotypes for Sepsis”  In JAMA 321.20, 2019, pp. 2003–2017  DOI: 10.1001/jama.2019.5791
  * [16] Veselin Stoyanov, Alexander Ropson and Jason Eisner  “Empirical Risk Minimization of Graphical Model Parameters Given Approximate Inference, Decoding, and Model Structure”  In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics JMLR WorkshopConference Proceedings, 2011, pp. 725–733  URL: https://proceedings.mlr.press/v15/stoyanov11a.html
  * [17] Naftali Tishby, Fernando C Pereira and William Bialek  “The information bottleneck method”  In arXiv preprint physics/0004057, 2000
  * [18] Annemarie M Wensing et al.  “2019 update of the drug resistance mutations in HIV-1”  In Topics in antiviral medicine 27.3 International Antiviral Society–USA, 2019, pp. 111
  * [19] Aleksander Wieczorek and Volker Roth  “On the Difference between the Information Bottleneck and the Deep Information Bottleneck”  In Entropy 22.2 MDPI AG, 2020, pp. 131  DOI: 10.3390/e22020131
  * [20] Maurizio Zazzi et al.  “Predicting response to antiretroviral treatment by machine learning: the EuResist project”  In Intervirology 55.2 Karger Publishers, 2012, pp. 123–127

Supplementary Materials

##  10 Derivations

In this section we provide the derivation of ELBO and coordinate ascent
updates for our prediction-focused models. The derivations shown below will be
for the pf-HMM, where $N$ is the number of sequences and $T_{n}$ signifies the
length of sequence $n$. The results for pf-GMMs directly follow by setting $N$
as the number of data points and setting each $T_{n}$ as 1. The parameters of
the model are $\Theta=\\{\theta,A,B,\pi,\eta\\}$

###  10.1 ELBO derivation

The lower bound for the log likelihood of one sequence with the posterior
distribution $q(\textbf{Z},\boldsymbol{\xi}|\Theta,\varphi)$ is:

$\displaystyle\log p(\textbf{X}=\textbf{x},\textbf{Y}=\textbf{y})$
$\displaystyle\geq\operatorname*{\mathbb{E}}_{q(\textbf{Z},\boldsymbol{\xi})}\left[\frac{\log
p(\textbf{Z},\boldsymbol{\xi},{\textbf{x},\textbf{y}}|\Theta,p)}{q(\textbf{Z},\boldsymbol{\xi}|\Theta,\varphi)}\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{q(\textbf{Z},\boldsymbol{\xi})}\left[\log
p({\textbf{x},\textbf{y}}|\textbf{Z},\boldsymbol{\xi},\Theta)\right]-\textbf{KL}\left(q(\textbf{Z},\boldsymbol{\xi}|\Theta,\varphi)||p(\textbf{Z},\boldsymbol{\xi}|A,{\theta},p)\right)$
$\displaystyle=\operatorname*{\mathbb{E}}_{q}\left[\log
p({\textbf{x}}|\textbf{Z},\boldsymbol{\xi},B,\pi)\right]+\operatorname*{\mathbb{E}}_{q}\left[\log
p({\textbf{y}}|\textbf{Z},\eta^{\mu},\eta^{V})\right]+\operatorname*{\mathbb{E}}_{q}\left[\log
p(\textbf{Z}|{\theta},A)\right]+\operatorname*{\mathbb{E}}_{q}\left[\log
p(\boldsymbol{\xi}|p)\right]$
$\displaystyle\quad-\operatorname*{\mathbb{E}}_{q}\left[\log
q(\boldsymbol{\xi}|\varphi)\right]-\operatorname*{\mathbb{E}}_{q}\left[\log
q(\textbf{Z}|{\textbf{x},\textbf{y}},\varphi,\Theta)\right]$

where the individual terms are (define
$q(Z_{t}|{\textbf{x},\textbf{y}},\Theta,\varphi)=q(Z_{t})$) :

  1. 1.

$\displaystyle\operatorname*{\mathbb{E}}_{q}\left[\log
p({\textbf{x}}|\textbf{Z},\boldsymbol{\xi},B,\pi)\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{q(\boldsymbol{\xi}|\varphi)q(\textbf{Z}|{\textbf{x},\textbf{y}},\Theta,\varphi)}\left[\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}{\xi_{t,d}^{n}}\log
p_{B}(x_{td}^{n}|{Z_{t}}^{n},B)+\left(1-{\xi_{t,d}^{n}}\right)\log
p_{\pi}(x_{td}^{n}|\pi)\right]$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}{\varphi_{d}}\operatorname*{\mathbb{E}}_{q(Z_{t}^{n})}\left[\log
p_{B}(x_{td}^{n}|{Z_{t}^{n}},B)\right]+\left(1-{\varphi_{d}}\right)\log
p_{\pi}(x_{td}^{n}|\pi)$

  2. 2.

$\displaystyle\operatorname*{\mathbb{E}}_{q}\left[\log
p({\textbf{y}}|\textbf{Z},\boldsymbol{\xi},B,\pi)\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{q(\textbf{Z}|{\textbf{x},\textbf{y}},\Theta,\varphi)}\left[\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\log
p(y_{t}^{n}|Z_{t}^{n},\eta^{\mu},\eta^{V})\right]=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\operatorname*{\mathbb{E}}_{q(Z_{t}^{n})}\left[\log
p(y_{t}^{n}|Z_{t}^{n},\eta^{\mu},\eta^{V})\right]$

  3. 3.

$\displaystyle\operatorname*{\mathbb{E}}_{q}\left[\log
p(\textbf{Z}|{\theta},A)\right]-\operatorname*{\mathbb{E}}_{q}\left[\log
q(\textbf{Z}|{\textbf{x},\textbf{y}},\varphi,\Theta)\right]$
$\displaystyle=\sum_{n=1}^{N}\operatorname*{\mathbb{E}}_{q(Z_{1}^{n})}\left[\log\frac{p(Z_{1}^{n}|{\theta})}{q(Z_{1}^{n}|{\textbf{x},\textbf{y}})}\right]+\sum_{n=1}^{N}\sum_{t=2}^{T_{n}}\operatorname*{\mathbb{E}}_{q(Z_{t-1,t}^{n}|{\textbf{x},\textbf{y}})}\left[\log\frac{p(Z_{t}^{n}|Z_{t-1}^{n},A)}{q(Z_{t}^{n}|Z_{t-1}^{n},{\textbf{x},\textbf{y}})}\right]$

  4. 4.

$\displaystyle\operatorname*{\mathbb{E}}_{q}\left[\log
p(\boldsymbol{\xi}|p)\right]-\operatorname*{\mathbb{E}}_{q}\left[\log
q(\boldsymbol{\xi}|\varphi)\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{q(\boldsymbol{\xi}|\varphi)}\left[\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}\log\frac{p(\xi_{t,d}^{n}|p)}{q(\xi_{t,d}^{n}|\varphi_{d})}\right]$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}\varphi_{d}\log\frac{p}{\varphi_{d}}+\left(1-\varphi_{d}\right)\log\frac{1-p}{1-\varphi_{d}}$

###  10.2 Coordinate Ascent Updates

We have the ELBO as:

$\displaystyle L$
$\displaystyle=\sum_{t}\left(\sum_{d=1}^{D}{\varphi_{d}}\operatorname*{\mathbb{E}}_{q(Z_{t}|{\textbf{x},\textbf{y}})}\left[\log
p_{B}(x_{td}|Z_{t})\right]+\left(1-\varphi_{d}\right)\log
p_{\pi}(x_{td})\right)+\operatorname*{\mathbb{E}}_{q(Z_{T}|{\textbf{x},\textbf{y}})}\left[\sum_{t=0}^{T_{n}}\log
p(y_{t}|Z_{t},\eta^{\mu},\eta^{V})\right]$
$\displaystyle-\sum_{d}\operatorname*{\mathbb{E}}_{q(\boldsymbol{\xi}_{d}|\varphi_{d})}\left[\log\frac{q(\boldsymbol{\xi}_{d}|\varphi_{d})}{p(\boldsymbol{\xi}_{d}|p_{d})}\right]-\operatorname*{\mathbb{E}}_{q(Z_{0}|{\textbf{x},\textbf{y}})}\left[\log\frac{q(Z_{0}|{\textbf{x},\textbf{y}})}{p(Z_{0}|{\theta})}\right]-\sum_{t}\operatorname*{\mathbb{E}}_{q(Z_{t-1,t}|{\textbf{x},\textbf{y}})}\left[\log\frac{q(Z_{t}|Z_{t-1},{\textbf{x},\textbf{y}})}{p(Z_{t}|Z_{t-1},A)}\right]$
$\displaystyle=\operatorname*{\mathbb{E}}_{q}\left[\log
p({\textbf{x}}|\textbf{Z},\boldsymbol{\xi},B,\pi)\right]+\operatorname*{\mathbb{E}}_{q}\left[\log
p({\textbf{y}}|\textbf{Z},\eta^{\mu},\eta^{V})\right]+\operatorname*{\mathbb{E}}_{q}\left[\log
p(\textbf{Z}|{\theta},A)\right]+\operatorname*{\mathbb{E}}_{q}\left[\log
p(\boldsymbol{\xi}|p)\right]$
$\displaystyle\quad-\operatorname*{\mathbb{E}}_{q}\left[\log
q(\boldsymbol{\xi}|\varphi)\right]-\operatorname*{\mathbb{E}}_{q}\left[\log
q(\textbf{Z}|{\textbf{x},\textbf{y}},\varphi,\Theta)\right]$

####  10.2.1 $\varphi$ update

Relevant terms:

$\displaystyle L_{\varphi}$
$\displaystyle=\operatorname*{\mathbb{E}}_{q}\left[\log
p({\textbf{x}}|\textbf{Z},\boldsymbol{\xi},B,\pi)\right]+\operatorname*{\mathbb{E}}_{q}\left[\log
p(\boldsymbol{\xi}|p)\right]-\operatorname*{\mathbb{E}}_{q}\left[\log
q(\boldsymbol{\xi}|\varphi)\right]$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}{\varphi_{d}}\operatorname*{\mathbb{E}}_{q(Z_{t})}\left[\log
p_{B}(x_{td}^{n}|{Z_{t}},B)\right]+\left(1-{\varphi_{d}}\right)\log
p_{\pi}(x_{td}^{n}|\pi)+\varphi_{d}\log\frac{p}{\varphi_{d}}+\left(1-\varphi_{d}\right)\log\frac{1-p}{1-\varphi_{d}}$

Gradient of the loss terms:

$\displaystyle\nabla_{\varphi_{d^{\prime}}}L_{\varphi}=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\operatorname*{\mathbb{E}}_{q(Z_{t})}\left[\log
p_{B}(x_{td^{\prime}}^{n}|{Z_{t}},B)\right]-\log
p_{\pi}(x_{td^{\prime}}^{n}|\pi)+\log\frac{p}{1-p}-\log\frac{\varphi_{d^{\prime}}}{1-\varphi_{d^{\prime}}}$

Update:

$\displaystyle\varphi_{d}$
$\displaystyle\leftarrow\sigma\left(\log\frac{p}{1-p}+\frac{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\operatorname*{\mathbb{E}}_{q(Z_{t}^{n}|{\textbf{x},\textbf{y}},\varphi)}\left[\log
p_{B}(x_{td}^{n}|Z_{t}^{n},{B}\right]-\log
p_{\pi}(x_{td}^{n}|{\pi})}{\sum_{n=1}^{N}T_{n}}\right)$

####  10.2.2 $B$ update

$B^{\mu}$

Relevant terms:

$\displaystyle L_{B^{\mu}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}{\varphi_{d}}\operatorname*{\mathbb{E}}_{q(Z_{t}^{n})}\left[\log
p_{B}(x_{td}^{n}|{Z_{t}^{n}},B)\right]$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}{\varphi_{d}}\operatorname*{\mathbb{E}}_{q(Z_{t}^{n})}\left[\frac{-1}{2B^{V}_{Z_{t}^{n},d}}\left(x_{td}^{n}-B^{\mu}_{Z_{t}^{n},d}\right)^{2}+c\right]$

Gradient of the loss terms:

$\displaystyle\nabla_{B^{\mu}_{j,d^{\prime}}}L_{B^{\mu}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}{\varphi_{d^{\prime}}}\nabla_{B^{\mu}_{j,d^{\prime}}}\sum_{k=1}^{K}q(Z_{t}^{n}=k)\cdot\frac{-1}{2B^{V}_{k,d^{\prime}}}\left(x_{td^{\prime}}^{n}-B^{\mu}_{k,d}\right)^{2}$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}{\varphi_{d^{\prime}}}q(Z_{t}^{n}=j)\cdot\frac{\left(x_{td^{\prime}}^{n}-B^{\mu}_{j,d^{\prime}}\right)}{B^{V}_{k,d^{\prime}}}$

Update:

$\displaystyle B^{\mu}_{j,d^{\prime}}$
$\displaystyle\leftarrow\frac{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}q(Z_{t}^{n}=j)x_{td^{\prime}}^{n}}{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}q(Z_{t}^{n}=j)}$

$B^{V}$

Relevant terms:

$\displaystyle L_{B^{V}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}{\varphi_{d}}\operatorname*{\mathbb{E}}_{q(Z_{t}^{n})}\left[-\frac{1}{2}\log
B^{V}_{Z_{t}^{n},d}-\frac{1}{2B^{V}_{Z_{t}^{n},d}}\left(x_{td}^{n}-B^{\mu}_{Z_{t}^{n},d^{\prime}}\right)^{2}+c\right]$

Gradient of the loss terms:

$\displaystyle\nabla_{B^{V}_{j,d^{\prime}}}L_{B^{V}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}{\varphi_{d^{\prime}}}q(Z_{t}^{n}=j)\left(-\frac{1}{2B^{V}_{j,d^{\prime}}}+\frac{1}{2\left(B^{V}_{j,d^{\prime}}\right)^{2}}\left(x_{td^{\prime}}^{n}-B^{\mu}_{j,d^{\prime}}\right)^{2}\right)$

Update (if $\varphi_{d^{\prime}}\neq 0$):

$\displaystyle B^{V}_{j,d^{\prime}}$
$\displaystyle\leftarrow\frac{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}q(Z_{t}^{n}=j)\left(x_{td^{\prime}}^{n}-B^{\mu}_{j,d^{\prime}}\right)^{2}}{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}q(Z_{t}^{n}=j)}$

####  10.2.3 $\pi$ update

$\pi^{\mu}$

Relevant terms:

$\displaystyle L_{\pi^{\mu}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}\left(1-{\varphi_{d}}\right)\log
p_{\pi}(x_{td}|\pi)$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}\left(1-{\varphi_{d}}\right)\frac{-1}{2\pi^{V}_{d}}\left(x_{td}^{n}-\pi^{\mu}_{d}\right)^{2}+c$

Gradient of the loss terms:

$\displaystyle\nabla_{\pi^{\mu}_{d^{\prime}}}L_{\pi^{\mu}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\left(1-\varphi_{d^{\prime}}\right)\cdot\frac{\left(x_{td^{\prime}}^{n}-\pi^{\mu}_{d^{\prime}}\right)}{\pi^{V}_{d^{\prime}}}$

Update (if $\varphi_{d^{\prime}}\neq 1$):

$\displaystyle\pi^{\mu}_{d^{\prime}}$
$\displaystyle\leftarrow\frac{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}x_{td^{\prime}}^{n}}{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}1}$
$\displaystyle=\frac{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}x_{td^{\prime}}^{n}}{\sum_{n=1}^{N}T_{n}}$

$\pi^{V}$

Relevant terms:

$\displaystyle L_{\pi^{V}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\sum_{d=1}^{D}\left(1-\varphi_{d}\right)\left(-\frac{1}{2}\log\pi^{V}_{d}-\frac{1}{2\pi^{V}_{d}}\left(x_{td}^{n}-\pi^{\mu}_{d^{\prime}}\right)^{2}+c\right)$

Gradient of the loss terms:

$\displaystyle\nabla_{\pi^{V}_{d^{\prime}}}L_{\pi^{V}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\left(1-\varphi_{d^{\prime}}\right)\left(-\frac{1}{2\pi^{V}_{d^{\prime}}}+\frac{1}{2\left(\pi^{V}_{d^{\prime}}\right)^{2}}\left(x_{td^{\prime}}^{n}-\pi^{\mu}_{d^{\prime}}\right)^{2}\right)$

Update:

$\displaystyle\pi^{V}_{d^{\prime}}$
$\displaystyle\leftarrow\frac{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}\left(x_{td^{\prime}}^{n}-\pi^{\mu}_{d^{\prime}}\right)^{2}}{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}1}$
$\displaystyle=\frac{\sum_{n=1}^{N}\sum_{t=1}^{T_{n}}{\left(x_{td^{\prime}}^{n}\right)^{2}}}{\sum_{n=1}^{N}T_{n}}-\left(\pi^{\mu}\right)^{2}$

####  10.2.4 ${\theta}$ update

Relevant terms (including the constraint and the regularizer):

$\displaystyle L_{{\theta}}$ $\displaystyle=\log
p({\theta}|\alpha)+\left\\{\sum_{n=1}^{N}\sum_{k=1}^{K}q(Z_{1}^{n}=k){\log
p(Z_{1}^{n}=k|{\theta})}\right\\}-\lambda_{{\theta}}\left(\sum_{k=1}^{K}{\theta}_{k}-1\right)$

Gradient of the loss terms:

$\displaystyle\nabla_{{\theta}_{k^{\prime}}}L_{{\theta}}$
$\displaystyle=-\lambda_{{\theta}}+(\alpha-1)\frac{1}{{\theta}_{k^{\prime}}}+\sum_{n=1}^{N}q(Z_{1}^{n}=k^{\prime})\frac{1}{{\theta}_{k^{\prime}}}$
$\displaystyle\nabla_{\lambda_{{\theta}}}L_{{\theta}}$
$\displaystyle=-\sum_{k=1}^{K}{\theta}_{k}$

Update:

$\displaystyle{\theta}_{k^{\prime}}$
$\displaystyle\leftarrow\frac{\alpha-1+\sum_{n=1}^{N}q(Z_{1}^{n}=k^{\prime})}{K\alpha-K+N}$
$\displaystyle=\frac{\alpha-1+\operatorname*{\mathbb{E}}\left[N^{1}_{k^{\prime}}\right]}{K\alpha-K+N}$

####  10.2.5 $A$ update

Relevant terms:

$\displaystyle L_{A}$
$\displaystyle=\left\\{\sum_{n=1}^{N}\sum_{t=2}^{T_{n}}\sum_{j=1}^{J}\sum_{k=1}^{K}q(Z_{t-1}^{n}=j,Z_{t}^{n}=k){\log
p(Z_{t}^{n}=k|Z_{t-1}^{n}=j,A)}\right\\}-\sum_{l=1}^{K}\lambda_{A_{l}}(\sum_{k=1}^{K}A_{lk}-1)$

Gradient of the loss terms:

$\displaystyle\nabla_{A_{j^{\prime},k^{\prime}}}L_{A}$
$\displaystyle=-\lambda_{A_{j^{\prime}}}+\sum_{n=1}^{N}\sum_{t=2}^{T_{n}}q(Z_{t-1}^{n}=j^{\prime},Z_{t}^{n}=k^{\prime})\frac{1}{A_{j^{\prime},k^{\prime}}}$
$\displaystyle\nabla_{\lambda_{A_{l^{\prime}}}}L_{A}$
$\displaystyle=-\sum_{k=1}^{K}A_{l^{\prime}k}$

Update:

$\displaystyle A_{j^{\prime},k^{\prime}}$
$\displaystyle\leftarrow\frac{\sum_{n=1}^{N}\sum_{t=2}^{T_{n}}q(Z_{t-1}^{n}=j^{\prime},Z_{t}^{n}=k^{\prime})}{\sum_{n=1}^{N}\sum_{t=2}^{T_{n}}\sum_{k=1}^{K}q(Z_{t-1}^{n}=j^{\prime},Z_{t}^{n}=k)}=\frac{\sum_{n=1}^{N}\sum_{t=2}^{T_{n}}q(Z_{t-1}^{n}=j^{\prime},Z_{t}^{n}=k^{\prime})}{\sum_{n=1}^{N}\sum_{t=2}^{T_{n}}q(Z_{t-1}^{n}=j^{\prime})}$
$\displaystyle=\frac{\operatorname*{\mathbb{E}}\left[N_{j^{\prime},k^{\prime}}\right]}{\sum_{k^{\prime}=1}^{K}\operatorname*{\mathbb{E}}\left[N_{j^{\prime},k^{\prime}}\right]}$

##  11 Proofs

In this section, we provide the proof for the limitation result in Section 6
of the main paper. Specifically, we prove that maximization of the joint
likelihood $\log p(\textbf{X},\textbf{Y})$ fails to give predictive components
when the irrelevant dimensions cannot be modeled well under the noise
distribution, $\pi$.

##### Proof

We have that the ELBO is equal to the log likelihood (for any parameters
$\Theta$) when
$KL(q(\textbf{Z},\boldsymbol{\xi})||p(\textbf{Z},\boldsymbol{\xi}|\textbf{X},\textbf{Y};{\Theta}))$
is zero:

$\displaystyle\log p(\textbf{X},\textbf{Y};\Theta)$
$\displaystyle=\operatorname*{\mathbb{E}}_{p(\textbf{Z},\boldsymbol{\xi}|\textbf{X},\textbf{Y};{\Theta})}\left[\log
p(\textbf{X},\textbf{Y},\textbf{Z},\boldsymbol{\xi};\Theta)\right]+\mathbb{H}\left[p(\textbf{Z},\boldsymbol{\xi}|\textbf{X},\textbf{Y};{\Theta})\right]$

Since the entropy term is bounded, we closely look at the first term:

$\displaystyle\operatorname*{\mathbb{E}}_{p(\textbf{Z},\boldsymbol{\xi}|\textbf{X},\textbf{Y};{\Theta})}\left[\log
p(\textbf{X},\textbf{Y},\textbf{Z},\boldsymbol{\xi};\Theta)\right]$
$\displaystyle=\sum_{d}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}}\left[\log
p(\xi_{d};p)\right]+\sum_{n}\operatorname*{\mathbb{E}}_{q_{Z_{n}}}\left[\log
p(Z_{n};\theta)\right]+\sum_{n}\operatorname*{\mathbb{E}}_{q_{Z_{n}}}\left[\log
p(Y_{n}|Z_{n};\eta)\right]$
$\displaystyle\quad+\sum_{n}\sum_{d}\operatorname*{\mathbb{E}}_{q_{\xi_{d},Z_{n}}}\left[\log
p(X_{nd}|\xi_{d},Z_{n};\Theta)\right]$
$\displaystyle=-\sum_{d}\mathbb{H}\left[p(\xi_{d};p)\right]-\sum_{n}\mathbb{H}\left[p(Z_{n};\theta)\right]+\sum_{n}\operatorname*{\mathbb{E}}_{q_{Z_{n}}}\left[\log
p(Y_{n}|Z_{n};\eta)\right]$
$\displaystyle\quad+\sum_{n}\sum_{d}\operatorname*{\mathbb{E}}_{q_{\xi_{d},Z_{n}}}\left[\log
p(X_{nd}|\xi_{d},Z_{n};\Theta)\right]$

where we use the short-hands:
$q_{\xi_{d}}:=p(\xi_{d}|\textbf{X},\textbf{Y};\Theta),q_{Z_{n}}:=p(Z_{n}|\textbf{X},\textbf{Y};\Theta),q_{\xi_{d},Z_{n}}:=p(\xi_{d},Z_{n}|\textbf{X},\textbf{Y};\Theta),q_{Z_{n}|\xi_{d}}:=p(Z_{n}|\xi_{d},\textbf{X},\textbf{Y};\Theta)$.
Again, we will generally have the first three terms bounded, we expand the
last term:

$\displaystyle\sum_{n}\sum_{d}\operatorname*{\mathbb{E}}_{q_{\xi_{d},Z_{n}}}\left[\log
p(X_{nd}|\xi_{d},Z_{n};\Theta)\right]$
$\displaystyle=\sum_{n}\sum_{d}\operatorname*{\mathbb{E}}_{q_{\xi_{d},Z_{n}}}\left[\xi_{d}\log
p(X_{nd}|\xi_{d}=1,Z_{n};B)+(1-\xi_{d})\log p(X_{nd}|\xi_{d}=0;\pi)\right]$
$\displaystyle=\sum_{n}\sum_{d}\operatorname*{\mathbb{E}}_{q_{\xi_{d},Z_{n}}}\left[\xi_{d}\log
p(X_{nd}|Z_{n};B_{d})+(1-\xi_{d})\log p(X_{nd};\pi_{d})\right]$
$\displaystyle=\sum_{n}\sum_{d}\log
p(X_{nd};\pi_{d})+\sum_{n}\sum_{d}\operatorname*{\mathbb{E}}_{q_{\xi_{d},Z_{n}}}\left[\xi_{d}\left(\log
p(X_{nd}|Z_{n};B_{d})-\log p(X_{nd};\pi_{d})\right)\right]$

$\displaystyle=\sum_{n}\sum_{d}\log
p(X_{nd};\pi_{d})+\sum_{n}\sum_{d}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}}\left[\xi_{d}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{d}}}\left[\log
p(X_{nd}|Z_{n};B_{d})\right]-\log p(X_{nd};\pi_{d})\right)\right]$ (4)

Now consider the two local optima where we either pick out the components
along $X_{1}$ or along $\textbf{X}_{2:6}$. Call them
$\Theta^{(1)}=\\{B^{(1)},\pi^{(1)},\theta^{(1)}\\}$ and
$\Theta^{(2)}=\\{B^{(2)},\pi^{(2)},\theta^{(2)}\\}$ respectively.

For $B^{(1)}$, the learned $B$ components would then be $\mathcal{N}(0,1)$ and
$\mathcal{N}(6,1)$ distributions for the first dimension and be the same as
$\pi^{(1)}_{2:6}$ for the remaining dimensions.

For $B^{(2)}$, the learned $B$ distribution would then be
$\mathcal{N}(\textbf{0},I_{5})$ and $\mathcal{N}(\bm{\mu},I_{5})$
distributions for dimensions 2-6, and the same as $\pi^{(2)}_{1}$ for the
first dimension.

Also note that $\pi^{(1)}=\pi^{(2)}$ because they are learned to just be the
best-fit univariate Gaussian per dimension (as derived in subsection 10.2.3).

Then for $\Theta^{(1)}$, the Eq. 4 becomes:

$\displaystyle=\sum_{n}\sum_{d}\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(1)})+\sum_{n}\sum_{d}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(1)}}\left[\xi_{d}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{d}}^{(1)}}\left[\log
p(X_{nd}|Z_{n};{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}B}_{d}^{(1)})\right]-\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(1)})\right)\right]$
$\displaystyle=\sum_{n}\sum_{d}\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(1)})+\sum_{n}\operatorname*{\mathbb{E}}_{q_{\xi_{1}}^{(1)}}\left[\xi_{1}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{1}}^{(1)}}\left[\log
p(X_{n1}|Z_{n};{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}B}_{1}^{(1)})\right]-\log
p(X_{n1};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{1}^{(1)})\right)\right]$
$\displaystyle\quad+\sum_{d=2}^{D}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(1)}}\left[\xi_{d}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{d}}^{(1)}}\left[\log
p(X_{nd}|Z_{n};{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}B}_{d}^{(1)})\right]-\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(1)})\right)\right]$

$\displaystyle=\sum_{n}\sum_{d}\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(1)})+\sum_{n}\operatorname*{\mathbb{E}}_{q_{\xi_{1}}^{(1)}}\left[\xi_{1}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{1}}^{(1)}}\left[\log
p(X_{n1}|Z_{n};{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}B}_{1}^{(1)})\right]-\log
p(X_{n1};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{1}^{(1)})\right)\right]$
$\displaystyle\quad+\sum_{d=2}^{D}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(1)}}\left[\xi_{d}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{d}}^{(1)}}\left[\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(1)})\right]-\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(1)})\right)\right]$

which finally reduces to:

$\displaystyle=\sum_{n}\sum_{d}\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(1)})+\sum_{n}\operatorname*{\mathbb{E}}_{q_{\xi_{1}}^{(1)}}\left[\xi_{1}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{1}}^{(1)}}\left[\log
p(X_{n1}|Z_{n};{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}B}_{1}^{(1)})\right]-\log
p(X_{n1};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{1}^{(1)})\right)\right]$
(5)

Similarly for $\Theta^{2}$, the Eq. 4 becomes:

$\displaystyle=\sum_{n}\sum_{d}\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(2)})+\sum_{n}\operatorname*{\mathbb{E}}_{q_{\xi_{1}}^{(2)}}\left[\xi_{1}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{1}}^{(2)}}\left[\log
p(X_{n1}|Z_{n};{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}B}_{1}^{(2)})\right]-\log
p(X_{n1};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{1}^{(2)})\right)\right]$
$\displaystyle\quad+\sum_{d=2}^{D}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(2)}}\left[\xi_{d}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{d}}^{(2)}}\left[\log
p(X_{nd}|Z_{n};{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}B}_{d}^{(2)})\right]-\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(2)})\right)\right]$
$\displaystyle=\sum_{n}\sum_{d}\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(2)})+\sum_{n}\operatorname*{\mathbb{E}}_{q_{\xi_{1}}^{(2)}}\left[\xi_{1}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{1}}^{(2)}}\left[\log
p(X_{n1};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{1}^{(2)})\right]-\log
p(X_{n1};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{1}^{(2)})\right)\right]$
$\displaystyle\quad+\sum_{d=2}^{D}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(2)}}\left[\xi_{d}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{d}}^{(2)}}\left[\log
p(X_{nd}|Z_{n};{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}B}_{d}^{(2)})\right]-\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(2)})\right)\right]$

which reduces to:

$\displaystyle=\sum_{n}\sum_{d}\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(2)})+\sum_{d=2}^{D}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(2)}}\left[\xi_{d}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{d}}^{(2)}}\left[\log
p(X_{nd}|Z_{n};{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}B}_{d}^{(2)})\right]-\log
p(X_{nd};{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\pi}_{d}^{(2)})\right)\right]$
(6)

Using Equations 5 and 6, we can write log likelihood ratio:

$\displaystyle\log p(\textbf{X},\textbf{Y};\Theta^{(1)})-\log
p(\textbf{X},\textbf{Y};\Theta^{(2)})$
$\displaystyle=\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\mathbb{H}\left[p(\textbf{Z},\boldsymbol{\xi}|\textbf{X},\textbf{Y};{\Theta^{(1)}})\right]-\mathbb{H}\left[p(\textbf{Z},\boldsymbol{\xi}|\textbf{X},\textbf{Y};{\Theta^{(2)}})\right]$
$\displaystyle\quad+\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\sum_{n}\mathbb{H}\left[p(Z_{n};\theta^{(2)})\right]-\sum_{n}\mathbb{H}\left[p(Z_{n};\theta^{(1)})\right]$
$\displaystyle\quad+\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\sum_{n}\operatorname*{\mathbb{E}}_{q_{Z_{n}}^{(1)}}\left[\log
p(Y_{n}|Z_{n};\eta^{(1)})\right]-\sum_{n}\operatorname*{\mathbb{E}}_{q_{Z_{n}}^{(2)}}\left[\log
p(Y_{n}|Z_{n};\eta^{(2)})\right]$
$\displaystyle\quad+\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\sum_{n}\operatorname*{\mathbb{E}}_{q_{\xi_{1}}^{(1)}}\left[\xi_{1}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{1}}^{(1)}}\left[\log
p(X_{n1}|Z_{n};B_{1}^{(1)})\right]-\log p(X_{n1};\pi_{1}^{(1)})\right)\right]$
$\displaystyle\quad-\sum_{d=2}^{D}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(2)}}\left[\xi_{d}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{d}}^{(2)}}\left[\log
p(X_{nd}|Z_{n};B_{d}^{(2)})\right]-\log p(X_{nd};\pi_{d}^{(2)})\right)\right]$
$\displaystyle=\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}O(1)$
$\displaystyle\quad+\color[rgb]{.75,0,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.75,0,.25}\sum_{n}\operatorname*{\mathbb{E}}_{q_{\xi_{1}}^{(1)}}\left[\xi_{1}\left(\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{1}}^{(1)}}\left[\log
p(X_{n1}|Z_{n};B_{1}^{(1)})\right]-\log p(X_{n1};\pi_{1}^{(1)})\right)\right]$
$\displaystyle\quad-\sum_{d=2}^{D}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(2)}}\left[\xi_{d}\operatorname*{\mathbb{E}}_{q_{Z_{n}|\xi_{d}}^{(2)}}\left[\log
p(X_{nd}|Z_{n};B_{d}^{(2)})\right]\right]+\sum_{d=2}^{D}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(2)}}\left[\xi_{d}\log
p(X_{nd};\pi_{d}^{(2)})\right]$
$\displaystyle=O(1)+\sum_{d=2}^{D}\operatorname*{\mathbb{E}}_{q_{\xi_{d}}^{(2)}}\left[\xi_{d}\log
p(X_{nd};\pi_{d}^{(2)})\right]$
$\displaystyle\rightarrow_{\mu\rightarrow\infty}-\infty$

which implies that the likelihood ratio goes to 0. Therefore, $\Theta^{(1)}$
will not be selected over $\Theta^{(2)}$ as $\mu\rightarrow\infty$.

##  12 Additional Experiments

###  12.1 Plot of the bound in Eq 3 in the main paper

Figure 6: Heldout AUROC vs pf-GMM hyperparameter $p$ when we optimize the
bound in Eq. 3 (main paper) instead of the likelihood—performance peaks for a
much larger range of p, giving it better task-focused properties over a larger
range of $p$. This would motivate looking for a tractable proxy for this
objective (which is left to future work).

###  12.2 pf-GMM Simulated Data Experiments

Figure 7: pf-GMM performance w.r.t. target AUROC on simulated data when we
vary the component budget for the mixture models. Each panel corresponds to a
dataset with a specified number of relevant dimensions, while keeping the
total number of dimensions fixed at 100. Notice our predictive performance is
close to optimal (as measured by a discriminative model e.g. Logistic
regression).

###  12.3 pf-HMM Simulated Data Experiments

Figure 8: pf-HMM performance w.r.t. target AUROC on simulated data when we
vary the component budget for the mixture models. Each panel corresponds to a
dataset with a specified number of relevant dimensions, while keeping the
total number of dimensions fixed at 20.

##  13 Experimental Details

###  13.1 Simulated datasets for the GMM experiments

We follow the following generative process to generate a collection of
datasets:

$\displaystyle p_{k}\sim\textbf{Unif}\left(\\{0.05,0.95\\}\right)\forall k$
$\displaystyle[\textbf{X}_{1},y]\sim\sum_{k=0}^{K-1}\theta_{k}^{\text{Rel}}[\mathcal{N}(K\cdot\textbf{6},I),\textbf{Bern}(p_{k})]$
$\displaystyle\textbf{X}_{2}\sim\sum_{k=0}^{K-1}\theta_{k}^{\text{Irrel}}\mathcal{N}(K\cdot\textbf{6},I)$
$\displaystyle\textbf{X}=[\textbf{X}_{1};\textbf{X}_{2}];\quad\textbf{X}_{1}\in{\mathbb{R}}^{D_{\text{Rel}}};\quad\textbf{X}_{2}\in{\mathbb{R}}^{D_{\text{Irrel}}}$

where $\textbf{Unif}\left(\\{a,b\\}\right)$ selects either $a$ or $b$
uniformly at random, $\textbf{Bern}(a)$ is a Bernoulli distribution with
parameter $a$, and 6 is a constant vector. This process ensures that the
target $y$ is correlated with the first $D_{\text{Rel}}$ dimensions of the
input X but not the last $D_{\text{Irrel}}$ dimensions. The gap between
successive components is fixed to 6 so that the clusters are well separated
and so that the relevant and irrelevant dimensions have equivalent
distributions a-priori.

In our experiments, we set:

$\displaystyle K=10$ $\displaystyle
D_{\text{Rel}}\in\\{10,20,30,40,50,80,100\\}$ $\displaystyle
D_{\text{Rel}}+D_{\text{Irrel}}=100$
$\displaystyle\theta^{\text{Rel}}=normalize(0.5+[0,\dots,K-1])$
$\displaystyle\theta^{\text{Irrel}}=normalize(1+[0,\dots,K-1])$

where $normalize(\textbf{x})=\textbf{x}/\|\textbf{x}\|_{1}$ makes sure the
vectors are valid probability distributions.

This gives us a knob to tweak: $D_{\text{Rel}}$ (number of relevant
dimensions). It is common for the input to only have a few relevant
dimensions. Therefore, we vary $D_{\textbf{rel}}$in { 10, 20, 30, 40, 50, 80,
100 }.

###  13.2 Simulated datasets for the HMM experiments

$\displaystyle[\textbf{X}_{1},y]\sim\sum_{k=0}^{K-1}\theta_{k}^{\text{Rel}}[\mathcal{N}(K\cdot\textbf{6},I),\textbf{Bern}(p_{k})]$
$\displaystyle\textbf{X}_{2}\sim\sum_{k=0}^{K-1}\theta_{k}^{\text{Irrel}}\mathcal{N}(K\cdot\textbf{6}K,I)$
$\displaystyle\textbf{X}=[\textbf{X}_{1};\textbf{X}_{2}];\quad\textbf{X}_{1}\in{\mathbb{R}}^{D_{\text{Rel}}};\quad\textbf{X}_{2}\in{\mathbb{R}}^{D_{\text{Irrel}}}$

In our experiments, we set:

$\displaystyle K=4$ $\displaystyle p=[0.05,0.95,0.05,0.95]$ $\displaystyle
D_{\text{Rel}}\in\\{2,3,4,5,8,12,20\\}$ $\displaystyle
D_{\text{Rel}}+D_{\text{Irrel}}=20$
$\displaystyle\theta^{\text{Rel}}=\theta^{\text{Irrel}}=[.1,.2,.3,.4]$
$\displaystyle A^{\text{Rel}}\leftarrow
normalize_{row}(.1+\begin{bmatrix}-z_{1}-\\\ \dots\\\
-z_{K}-\end{bmatrix}+\textbf{I}_{K})$ $\displaystyle
A^{\text{Irrel}}\leftarrow
normalize_{row}(.01+\begin{bmatrix}-z^{\prime}_{1}-\\\ \dots\\\
-z^{\prime}_{K}-\end{bmatrix}+\textbf{I}_{K})$ $\displaystyle
z_{i},z^{\prime}_{i}\sim\textbf{Cat}(\textbf{1}/K),z_{i},z^{\prime}_{i}\in{\mathbb{R}}^{K}$

where $normalize_{row}(M)$ applies $normalize$ to each row of matrix $M$.
Also, the numbers .1 and .01 are arbitrary choices to differentiate
$A^{\text{Rel}}$ and $A^{\text{Irrel}}$, and the simulation would work with
different values too.

###  13.3 HIV Experiments

Therapy for HIV involves administering cocktails of antiretrovirals from five
classes namely, Non-nucleoside Reverse Transcriptase Inhibitors (nnRTIs),
Nucleoside Reverse Transcriptase Inhibitors (nRTIs), Protease Inhibitors
(PIs), Fusion Inhibitors (FIs), and Integrase Inhibitors (IIs) to bring the
viral load below detection limits ($\leq 40$ copies/ml). We study 53 236
patients with HIV from the EuResist Integrated Database. Each person has a
time-series of average length 16 steps where a time step is approximately 4
months between consecutive treatments. Our task is to predict whether a
treatment will bring the viral load below detection limits in the next time-
step. Each input contains 138 features including CD4+counts, genetic
mutations, treatments in terms of drug classes and lab results. Though it is
common to have many genetic mutations, only a few of these may be relevant for
inducing drug resistance thus increasing the viral load.