SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ethan Weinberger
Paul G. Allen School of Computer Science
Seattle, WA 98195
ewein@cs.washington.edu
Ian C. Covert
icovert@stanford.edu
Su-In Lee
suinlee@cs.washington.edu
Contrastive analysis (CA) refers to the exploration of variations uniquely enriched in a target dataset as compared to a corresponding background dataset generated from sources of variation that are irrelevant to a given task. For example, a biomedical data analyst may wish to find a small set of genes to use as a proxy for variations in genomic data only present among patients with a given disease (target) as opposed to healthy control subjects (background). However, as of yet the problem of feature selection in the CA setting has received little attention from the machine learning community. In this work we present contrastive feature selection (CFS), a method for performing feature selection in the CA setting. We motivate our approach with a novel information-theoretic analysis of representation learning in the CA setting, and we empirically validate CFS on a semi-synthetic dataset and four real-world biomedical datasets. We find that our method consistently outperforms previously proposed state-of-the-art supervised and fully unsupervised feature selection methods not designed for the CA setting. An open-source implementation of our method is available at https://github.com/suinleelab/CFS.
In many scientific domains, it is increasingly common to analyze high-dimensional datasets consisting of observations with many features [1, 2, 3]. However, not all features are created equal: noisy, information-poor features can obscure important structures in a dataset and impair the performance of machine learning algorithms on downstream tasks. As such, it is often desired to select a small number of informative features to consider in subsequent analyses. For example, biologists may seek to measure the expression levels of a small number of genes and use these in a variety of future prediction tasks, such as disease subtype prediction or cell type classification. Measuring only a subset of features may also yield other benefits, including reduced experimental costs, enhanced interpretability, and better generalization for models trained with limited data.
While feature selection is relatively straightforward when supervision guides the process (e.g., class labels), such information is often unavailable in practice. In the unsupervised setting, we can select features for specific objectives like preserving clusters [4, 5] or local structure in the data [6, 7], or we can aim to simply reconstruct the full feature vector [8, 3]. However, these approaches are likely to focus on factors that dominate in the dataset, even if these structures are uninteresting for a specific
downstream analysis. A third setting that lies between supervised and unsupervised is contrastive analysis [9, 10, 11], which we describe below.
In the contrastive analysis (CA) setting, we have two datasets conventionally referred to as the target and background [10]. The target dataset contains observations with variability from both "interesting" (e.g., due to a specific disease under study) and "uninteresting" (e.g., demographic-related) sources. On the other hand, the corresponding background dataset (e.g., measurements from healthy patients with similar demographic profiles) is assumed to only contain variations due to the uninteresting sources. The goal of CA is then to isolate the patterns enriched in the target dataset and not present in the background for further study. Here, the target-background dataset pairing provides a form of weak supervision to guide our analysis, and carefully exploiting their differences can help us focus on features that are maximally informative of the target-specific variations.
For example, suppose our goal were to measure gene expression in cancer patients for a subset of genes to better understand molecular subtypes of cancer. Identifying factors specific to cancerous tissues is of great importance to the cancer research community, as these intra-cancer variations are essential to study disease progression and treatment response [12]. However, selecting an appropriate set of genes is not straightforward, as unsupervised methods are likely to be confounded by variations due to nuisance factors such as age, ethnicity, or batch effects, which often dominate the overall variation in gene expression and other omics datasets (genomics, transcriptomics, proteomics) [10, 13, 14]. Moreover, fully supervised methods may not be applicable as they require detailed labels that may not be available a priori (e.g., due to the rapid rate of development and testing of new cancer treatments [15]). Thus, methods that can effectively leverage the weak supervision available in CA may be able to select more appropriate features for a given analysis.
Isolating salient variations is the focus of many recent works on CA, including many that generalize unsupervised representation learning methods to the CA setting [9, 10, 11, 16, 17, 18, 19, 13]. This work aims to develop a principled approach for contrastive feature selection, which represents an important use case not considered in prior work.
The remainder of this manuscript is organized as follows: we begin (Section 2) with an introduction to the contrastive analysis setting and the unique challenges therein, followed by a discussion of related work (Section 3). We then proceed propose CFS, a feature selection method that leverages the weak supervision available in the CA setting (Section 4). CFS selects features that reflect variations enriched in a target dataset but are not present in the background data, and it does so using a novel two-step procedure that optimizes the feature subset within a neural network module. Subsequently, we develop a novel information-theoretic characterization of representation learning in the CA setting to justify our proposed method (Section 5); under mild assumptions about the data generating process, we prove that our two-step learning approach maximizes the mutual information with target-specific variations in the data. Finally, we validate our approach empirically through extensive experiments on a semi-synthetic dataset introduced in prior work as well as four real-world biomedical datasets (Section 6). We find that CFS consistently outperforms supervised and unsupervised feature selection algorithms, including state-of-the-art methods that it builds upon.
We now describe the problem of feature selection in the CA setting. Our analysis focuses on an observed variable  $\pmb{x} \in \mathbb{R}^d$  that we assume depends on two latent sources of variation: a salient variable  $s$  that represents interesting variations (e.g., disease subtype or progression), and a separate background variable  $z$  that represents variations that are not relevant for the intended analysis (e.g., gender or demographic information). Thus, we assume that observations are drawn from the data distribution  $p(\pmb{x}) = \int p(\pmb{x} \mid \pmb{z}, \pmb{s}) p(\pmb{z}, \pmb{s}) dz ds$ , which we depict using a graphical model in Figure 1.
Our goal is to select a small number of features that represent the relevant variability in the data, or that contain as much information as possible about  $s$ . To represent this, we let  $S \subseteq [d] \equiv \{1, \dots, d\}$  denote a subset of indices and  $x_{S} \equiv \{x_{i} : i \in S\}$  a subset of features. For a given value  $k < d$ , we aim to find  $S$  with  $|S| = k$  where  $x_{S}$  is roughly equivalent to  $s$ , or is similarly useful for downstream
<sup>1</sup>We note that contrastive analysis (CA) is unrelated to contrastive learning in the sense of SimCLR [20], CLIP [21], or other self-supervised learning methods [22]. We retain the name used in prior work [9, 10, 16, 17].
Figure 1: Data generating process for target and background samples. Observed values  $x$  are generated from two latent variables: a salient variable  $s$  and a background variable  $z$ . Both are active in the target dataset, while the salient variable is fixed at  $s = s'$  in the background data.
tasks (e.g., clustering cells by disease status). We state this goal intuitively here following prior work [11], but we later frame our objective through the lens of mutual information (Section 5).
What distinguishes CA from the supervised setting is that we lack labels to guide the feature selection process [23]. We could instead perform unsupervised feature selection [3], but this biases our selections towards variations that dominate the data generating process; thus, in real-world genomics datasets, we may inadvertently focus on demographic features that are irrelevant to a specific disease. Instead, the CA setting provides a form of weak supervision via two separate unlabeled datasets [11]. First, we have a target dataset  $\mathcal{D}_t$  consisting of samples with both sources of variation, or which follow the data distribution  $p(\boldsymbol{x})$  defined above (Figure 1). Next, we have a background dataset  $\mathcal{D}_b$  that consists of samples with no salient variation; this may represent a single fixed disease status, or perhaps no disease, and we formalize these as samples from  $p(\boldsymbol{x} \mid \boldsymbol{s} = s')$  for a fixed value  $s'$ . We note that individual samples from the target and background datasets are not paired and we do not assume the same number of samples from the target and background datasets (i.e.,  $|\mathcal{D}_t| \neq |\mathcal{D}_b|$ ). By analyzing both datasets, we hope to distinguish sources of variability specific to  $s$  rather than  $z$ .
The challenge we must consider is how to properly exploit the weak supervision provided by separate target and background datasets. To identify a maximally informative feature set  $\boldsymbol{x}_S$ , we must effectively disentangle the distinct sources of variation in the data, and we wish to do so while allowing for different dataset sizes  $|\mathcal{D}_b| = m$  and  $|\mathcal{D}_t| = n$ , greater variability in  $\boldsymbol{x}$  driven by irrelevant information  $\boldsymbol{z}$  than salient factors  $\boldsymbol{s}$ , and with no need for additional labeling. Our work does so by building on previous work on contrastive representation learning [10, 16, 18] and neural feature selection [3, 23], while deriving insights that apply beyond the feature selection problem.
Contrastive analysis Previous work has designed a number of representation learning methods to exploit the weak supervision available in the CA setting. Zou et al. [9] initially adapted mixture models for CA, and recent work adapted PCA to summarize target-specific variations while ignoring uninteresting factors of variability [10]. Other recent works developed probabilistic latent variable models [18, 16, 17, 19], including a method based on VAEs [11]. Severson et al. [16] experimented with adding a sparsity prior to perform feature selection within their model, but the penalty was designed for fully supervised scenarios: the penalty was manually tuned by observing how well pre-labeled target samples separated in the resulting latent space, which is not suitable for the scenario described in Section 2 where such labels are unavailable. Our work is thus the first to focus on feature selection in this setting, and we provide an information-theoretic analysis for our approach that generalizes to other forms of representation learning in the CA setting.
Feature selection Feature selection has been extensively studied in the machine learning literature, see [24, 25] for recent reviews. Such methods can be divided broadly into filter [26, 27], wrapper [28, 29] and embedded [30, 31] methods, and they are also distinguished by whether they require labels (i.e., supervised versus unsupervised). Recent work has introduced mechanisms for neural feature selection [32, 3, 23, 33, 34, 35], which learn a feature subset jointly with an expressive model and allow a flexible choice of prediction target and loss function. Our work builds upon these methods, which represent the current state-of-the-art in feature selection. However, unlike prior works that
Figure 2: Overview of contrastive feature selection (CFS). Our approach consists of a two-step optimization procedure. In the first step (left), we train an autoencoder on a background dataset, which we assume is generated solely from irrelevant variations due to  $z$ . In the second step (right), we train a feature selector layer using our target dataset. To encourage the selection of features that best reflect salient variations due to  $s$ , the selector layer's output is concatenated with the output of the fixed encoder  $g$ , and then passed to a learned reconstruction function  $f$ . Ideally, the feature set  $x_{S}$  chosen by the selector will capture only the factors of variation unique to the target data.
focus on either supervised [23, 34] or unsupervised feature selection [3, 36], we focus on contrastive dataset pairs and consider how to properly exploit the weak supervision in this setting.
4 Proposed method: contrastive feature selection
We aim to select  $k$  features within  $x \in \mathbb{R}^d$  that provide maximum information about the salient latent variable  $s$  in our data generating process (Section 2). Such a procedure must be performed carefully, as standard unsupervised methods may optimize for selecting features that capture much of the overall variance in a dataset but not the (potentially subtle) variations of interest due to  $s$ .
To do so, we begin with the following intuitive idea: suppose we have a priori a variable  $\pmb{b} \in \mathbb{R}^l$  with  $l < d$  that serves as a proxy for  $z$  (i.e., it summarizes all the uninteresting variation in the data). We assume for now that  $\pmb{b}$  is already available, and that it should explain a significant portion of  $x$ 's variance, so that any remaining variance is due solely to factors of interest. We can then aim to find a complementary set of features that capture the remaining variations in  $x$  not explained by  $\pmb{b}$ , and we do so by learning the feature set in conjunction with a function that reconstructs the full vector  $x$ . The reconstruction function  $f\colon \mathbb{R}^{k + l} \mapsto \mathbb{R}^d$  uses both  $\pmb{b}$  and the learned feature set  $x_S$  as inputs, and we can learn  $f$  and  $S$  jointly as follows:
$$
\underset {\theta , | S | = k} {\arg \min } \mathbb {E} _ {\boldsymbol {x} \sim \mathcal {D} _ {t}} \left| \left| f (\boldsymbol {b}, \boldsymbol {x} _ {S}; \theta) - \boldsymbol {x} \right| \right| ^ {2}. \tag {1}
$$
We call this approach contrastive feature selection (CFS). Notably, eq. (1) requires data with both background and salient variations, and thus it uses only the target dataset  $\mathcal{D}_t$ . Intuitively, optimizing this objective will select the set of  $k$  features  $S$  that best capture the variations present in the target dataset but which are not captured by the variable  $b$ .
To implement CFS we must make two choices: (1) how to optimize the feature set  $S \subseteq [d]$ , and (2) how to obtain  $b$ . First, assuming we already have  $b$ , we can leverage recent techniques for differentiable feature selection to learn  $S$  [23, 3]. Such methods select features by optimizing continuous approximations of discrete variables at the input layer of a neural network: selected variables are multiplied with a 1 and allowed to pass through the gates, while rejected features are multiplied with a 0. In our implementation of CFS, we use the Gaussian-based relaxation of Bernoulli
random variables termed stochastic gates (STG), proposed by Yamada et al. [23].<sup>2</sup> Let  $G \in [0,1]^d$  denote a random vector of STG values parameterized by  $\mu \in \mathbb{R}^d$ , so that samples of  $G$  are defined as
$$
\boldsymbol {G} _ {i} = \max  (0, \min  (1, \mu_ {i} + \zeta)), \tag {2}
$$
where  $\mu_{i}$  is a learnable parameter and  $\zeta$  is drawn from  $\mathcal{N}(0,\sigma^2)$ . By multiplying each feature by its corresponding gate value  $G_{i}$  and incorporating a regularization term that penalizes open gates (i.e., the number of indices with  $\mu_{i} > 0$ ), feature selection can be performed in a fully differentiable manner. Leveraging the STG mechanism, we obtain the following continuous relaxation of eq. (1):
$$
\underset {\boldsymbol {\mu}, \theta} {\arg \min } \mathbb {E} _ {\boldsymbol {x} \sim \mathcal {D} _ {t}} | | f (\boldsymbol {b}, \boldsymbol {x} \odot \boldsymbol {G}; \theta) - \boldsymbol {x} | | ^ {2} + \lambda \sum_ {i = 1} ^ {d} \Phi \left(\frac {\mu_ {i}}{\sigma}\right), \tag {3}
$$
where  $\odot$  denotes the Hadamard product,  $\lambda$  is a hyperparameter than controls the number of selected features, and  $\Phi$  is the standard Gaussian CDF [23].
Next, we must decide how to obtain  $\pmb{b}$  so that it captures irrelevant variations due to  $\pmb{z}$ , and ideally not those due to the salient variable  $\pmb{s}$ . We propose to learn it as a function of  $\pmb{x}$ , so that  $\pmb{b} = g(\pmb{x};\phi)$  and  $g:\mathbb{R}^d\mapsto \mathbb{R}^l$ . Intuitively, this is where we can leverage the background dataset, because it contains only uninteresting variations due to  $\pmb{z}$ . We learn this background representation function in conjunction with a reconstruction network  $h:\mathbb{R}^l\mapsto \mathbb{R}^d$  using the following objective:
$$
\underset {\phi , \eta} {\arg \min } \mathbb {E} _ {\boldsymbol {x} \sim \mathcal {D} _ {b}} | | h (g (\boldsymbol {x}; \phi); \eta) - \boldsymbol {x} | | ^ {2}. \tag {4}
$$
Previous representation learning methods for CA, such as the contrastive VAEs of Abid and Zou [11] and Severson et al. [16], also considered learning two representations; however, they learned both representations jointly by training simultaneously with target and background data. In our feature selection context, this joint training setup would entail optimizing with the sum of eq. (3) and eq. (4), so that  $g(\pmb{x};\phi)$  is updated based on both  $\mathcal{D}_b$  and  $\mathcal{D}_t$ . This is potentially problematic, because the background representation  $\pmb{b} = g(\pmb{x};\phi)$  may absorb salient information that then cannot be learned by  $\pmb{x}_S$ . Indeed, subsequent work observed that such joint training procedures may fail to properly segregate shared and target-specific variations, with target-specific variations being captured by the background representation function and vice versa [37].
In order to avoid this undesirable "leakage" of salient information into our background representation, we thus separate our objectives into a two-stage optimization procedure. For a given target-background dataset pair, we use the first stage to train an autoencoder network using only background samples (eq. (4)). In the second stage, we then use the encoder network  $\pmb{b} = g(\pmb{x};\phi)$  as our background proxy, freeze its weights, and learn a feature set  $\pmb{x}_S$  by optimizing eq. (3) using only target samples. We depict this procedure graphically in Figure 2. In the following section, we further motivate this two-stage approach by developing an information-theoretic formulation of representation learning in the CA setting, which illustrates why this approach compares favorably to joint contrastive training [11, 16], as well as fully unsupervised feature selection approaches [3].
5 Information-theoretic analysis
In this section, we motivate our approach from Section 4 through a framework of mutual information maximization. To justify learning two representations that are optimized in separate steps, we frame each step as maximizing a specific information-theoretic objective, which we then relate to our ultimate goal of representing the salient variation  $s$  via a feature set  $x_{S}$ .
We simplify our analysis by considering that all variables  $(\boldsymbol{x}, s, z)$  are all discrete rather than continuous, which guarantees that they have non-negative entropy [38]. The learned representations are deterministic functions of  $\boldsymbol{x}$ , and we generalize our analysis by letting the representation from the second step be an arbitrary variable  $\boldsymbol{a}$  rather than a feature set  $\boldsymbol{x}_S$ . Thus, our analysis applies even beyond the feature selection setting.
2While we focused on the STG layer in this work, CFS can be implemented with other choices of selection layers (e.g. the concrete selector layer of [3]). We refer the reader to Appendix A for further details.
5.1 Analyzing our approach
Before presenting our main results, we state two sets of assumptions required for our analysis. Regarding the data generating process (see Figure 1), we require mild conditions about the latent variables being independent and about each variable being explained by the others.
Assumption 1. (Assumptions on data distribution.) We assume that the latent variables  $s, z$  are roughly independent, that  $x$  is well-explained by the latent variables, and that each latent variable is well-explained by  $x$  and the other latent variable. That is, we assume a small constant  $\epsilon > 0$  such that  $I(s; z) \leq \epsilon$ ,  $H(x \mid s, z) \leq \epsilon$ ,  $H(s \mid x, z) \leq \epsilon$ , and  $H(z \mid x, s) \leq \epsilon$ .
Intuitively, we require independence so that salient information is not learned in the first step with  $\pmb{b}$  and then ignored by  $\pmb{a}$ ; the other conditions guarantee that  $s$  can be recovered at all via  $\pmb{x}$ . We also require an assumption about the learned representations, which is that they do not affect the independence between the latent variables  $s, z$ . This condition is not explicitly guaranteed by our approach, but it also is not discouraged, so we view it as a mild assumption.
Assumption 2. (Assumptions on learned representations.) We assume that  $s, z$  remain roughly independent after conditioning on the learned representations  $a, b$ . That is, we assume a small constant  $\epsilon > 0$  such that  $I(s; z \mid b) \leq \epsilon$  and  $I(s; z \mid a) \leq \epsilon$ .
With these assumptions, we can now present our analysis. First, we re-frame each step of our approach in an information-theoretic fashion:
- The first step in Section 4 is analogous to maximizing the mutual information  $I(\pmb{b}; \pmb{x} \mid s = s')$ . We do this for only a single value  $s'$  in practice, but for our analysis we consider that this is similar to  $I(\pmb{b}; \pmb{x} \mid s)$ , which would be achieved if we optimized  $\pmb{b}$  over multiple background datasets with different values. This holds because  $I(\pmb{b}; \pmb{x} \mid s) = \mathbb{E}_s[I(\pmb{b}; \pmb{x} \mid s = s)]$ .
- The second step in Section 4 is analogous to maximizing the mutual information  $I(\pmb{a}, \pmb{b}; \pmb{x})$ . However, by treating  $\pmb{b}$  as fixed and rewriting the objective as  $I(\pmb{a}, \pmb{b}; \pmb{x}) = I(\pmb{a}; \pmb{x} \mid \pmb{b}) + I(\pmb{b}; \pmb{x})$ , we see that we are in fact only maximizing  $I(\pmb{a}; \pmb{x} \mid \pmb{b})$ , or encouraging  $\pmb{a}$  to complement  $\pmb{b}$ .
- Finally, our true objective is for  $\mathbf{a}$  to be equivalent to  $s$ , which we view as maximizing  $I(\mathbf{a}; s)$ . If their mutual information is large, then  $\mathbf{a}$  is as useful as  $s$  in downstream prediction tasks [38]. The challenge in CA is optimizing this when we lack labeled training data.
Now, we have the following result that shows our two-step procedure implicitly maximizes  $I(\pmb{a}; \pmb{s})$  which is our true objective in contrastive analysis (proof in Appendix B).
Theorem 1. Under Assumptions 1 and 2, learning  $\mathbf{b}$  by maximizing  $I(\mathbf{b};\mathbf{x}\mid \mathbf{s})$  and learning  $\mathbf{a}$  by maximizing  $I(\mathbf{a};\mathbf{x}\mid \mathbf{b})$  yields the following lower bound on the mutual information  $I(\mathbf{a};\mathbf{s})$ :
$$
I (\boldsymbol {a}; \boldsymbol {x} \mid \boldsymbol {b}) + I (\boldsymbol {b}; \boldsymbol {x} \mid \boldsymbol {s}) - H (\boldsymbol {z}) - 4 \epsilon \leq I (\boldsymbol {a}; \boldsymbol {s}). \tag {5}
$$