SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The result above shows that our true objective is related to  $I(\pmb{b}; \pmb{x} \mid s)$  and  $I(\pmb{a}; \pmb{x} \mid \pmb{b})$ , which represent the two steps of our learning approach. Specifically, the sum of these terms provides a lower bound on  $I(\pmb{a}; s)$ , so maximizing them implicitly maximizes  $I(\pmb{a}; s)$  (similar to the ELBO in variational inference [39]). Within the inequality,  $H(z)$  acts as a rough upper limit on  $I(\pmb{b}; \pmb{x} \mid s)$ , and the constant  $\epsilon$  from Assumptions 1 and 2 contributes to looseness in the bound, which becomes tight in the case where  $\epsilon = 0$ . This result requires several steps to show, but it is perhaps intuitive that letting  $\pmb{b}$  serve as a proxy for  $z$  and then learning  $\pmb{a}$  as a complementary variable helps recover the remaining variation, or encourages a large value for  $I(\pmb{a}; s)$ .
5.2 Comparison to other approaches
We now compare our result for our two-stage approach with two alternatives: (1) a simpler one-stage "joint" training approach similar to previously proposed contrastive VAEs [11, 16], and (2) an unsupervised learning approach similar to the concrete autoencoder of Balin et al. [3].
First, the joint training approach learns  $\mathbf{a}$  and  $\mathbf{b}$  simultaneously rather than separately, which we can view as optimizing one objective corresponding to each dataset (background and target). Following the argument in Section 5.1, the background data is used roughly to maximize  $I(\mathbf{b};\mathbf{x}\mid \mathbf{s})$ . However,
the target data is used to maximize  $I(\mathbf{a}, \mathbf{b}; \mathbf{x})$  with respect to both  $\mathbf{a}$  and  $\mathbf{b}$  rather than just  $\mathbf{a}$ . Thus, we can modify the inequality in eq. (5) to obtain the following lower bound on  $I(\mathbf{a}; \mathbf{s})$ :
$$
I (\boldsymbol {a}, \boldsymbol {b}; \boldsymbol {x}) + I (\boldsymbol {b}; \boldsymbol {x} \mid \boldsymbol {s}) - \underbrace {I (\boldsymbol {b} ; \boldsymbol {x})} _ {\text {M a x i m i z e d}} - H (\boldsymbol {z}) - 4 \epsilon \leq I (\boldsymbol {a}; \boldsymbol {s}). \tag {6}
$$
This follows directly from Theorem 1 and the definition of mutual information [38]. By maximizing  $I(a, b; x) = I(b; x) + I(a; x \mid b)$  rather than just  $I(a; x \mid b)$ , we may inadvertently reduce the lower bound through  $I(b; x)$  and therefore limit the salient information that  $a$  is forced to recover. Indeed, our experiments show that the joint training approach results in lower-quality features compared to our proposed two-stage method (Section 6).
Next, we consider fully unsupervised methods that learn a single representation  $\mathbf{a}$ . Such methods do not leverage paired contrastive datasets, so we formalize them as maximizing the mutual information  $I(\mathbf{a};\mathbf{x})$ , which is natural when  $\mathbf{a}$  is used to reconstruct  $\mathbf{x}$  [3, 40]. Intuitively, if most variability in the data is due to irrelevant variation in  $\mathbf{z}$ , these methods are encouraged to maximize  $I(\mathbf{a};\mathbf{z})$  rather than  $I(\mathbf{a};\mathbf{s})$ . This can be seen through the following result (see proof in Appendix B).
Theorem 2. When we learn a single representation  $\mathbf{a}$  by maximizing  $I(\mathbf{a};\mathbf{x})$ , we obtain the following simultaneous lower bounds on  $I(\mathbf{a};\mathbf{s})$  and  $I(\mathbf{a};\mathbf{z})$ :
$$
I (\boldsymbol {a}; \boldsymbol {x}) - H (\boldsymbol {x}) + I (\boldsymbol {x}; \boldsymbol {s}) \leq I (\boldsymbol {a}; \boldsymbol {s})
$$
$$
I (\boldsymbol {a}; \boldsymbol {x}) - H (\boldsymbol {x}) + I (\boldsymbol {x}; \boldsymbol {z}) \leq I (\boldsymbol {a}; \boldsymbol {z}). \tag {7}
$$
The implication of Theorem 2 is that if  $z$  explains more variation than  $s$  in the observed variable  $x$ , or if we have  $I(\pmb{x};z) > I(\pmb{x};s)$ , then we obtain a higher lower bound for  $I(\pmb{a};z)$  than  $I(\pmb{a};s)$ . This is undesirable, because the mutual information with  $x$  roughly decomposes between  $s$  and  $z$ : as we show in Appendix B, Assumptions 1 and 2 imply that we have  $|I(\pmb{a};s) + I(\pmb{a};z) - I(\pmb{a};\pmb{x})| \leq 2\epsilon$ .
We conclude that under this information-theoretic perspective and several mild assumptions, we can identify concrete benefits to a two-stage approach that leverages paired contrastive datasets.
5.3 Caveats and practical considerations
There are several practical points to discuss about our analysis, but we first emphasize the aim of this section: our goal is to justify a two-stage approach that leverages paired contrastive datasets, and to motivate performing the optimization in these steps separately rather than jointly. Our approach provides clear practical benefits (Section 6), and this analysis provides a lens to understand why that is, even if it does not perfectly describe how the algorithm works in practice.
The caveats to discuss are the following. (1) In practice, we minimize mean squared error for continuous variables rather than maximizing mutual information; however, the two are related, and Appendix C shows that minimizing mean squared error maximizes a lower bound on the mutual information. (2) Our analysis relates to an arbitrary variable  $\mathbf{a}$  rather than a feature set  $\mathbf{x}_S$ . This is not an issue, and in fact implies that our results apply to more general representation learning and could improve existing methods [11]. (3) The assumptions we outlined may not hold in practice (see Assumptions 1 and 2). However, this represents a legitimate failure mode rather than an issue in our analysis: for example, if the latent variables are highly dependent, the second step may not recover meaningful variation from  $s$  because this will be captured to some extent by  $\mathbf{b}$  in the first step.
Here, we empirically validate our proposed CFS approach by using it to select maximally informative features in multiple target-background dataset pairs that were analyzed in previous work. We begin by considering a semi-synthetic dataset, Grassy MNIST [10], which lets us control the data generation process and the strength of uninteresting variations; we then report results on four publicly available biomedical datasets. For each target-background dataset pair, samples in the target dataset have ground-truth class labels related to the salient variations, which we use to evaluate the selected features. However, distinguishing these classes may be difficult using methods not designed for CA.
Figure 3: a-b, Example target (a) and background (b) images from the Grassy MNIST dataset. Our goal here is to select the set of features that best capture the target-dataset-specific variations due to the MNIST digits. c-d,  $k = 20$  features selected by CFS (c) versus unsupervised (CAE [3], DUFS [36]) and supervised (STG [23], LassoNet [34]) baselines (d).
To illustrate the benefits of leveraging the weak supervision available in the CA setting, we compared against unsupervised feature selection methods applied to our target datasets. In particular, we considered two recently proposed methods that claim state-of-the-art results in fully unsupervised settings: the concrete autoencoder (CAE) [3] and differentiable unsupervised feature selection (DUFS) [36]. Moreover, to illustrate that naive adaptations of previous methods designed for fully supervised settings are not sufficient to fully exploit the weak supervision available in the CA setting, we compared against the state-of-the-art fully supervised methods STG [23] and LassoNet [34] trained to simply classify samples as target versus background.
Finally, to illustrate the importance of carefully learning CFS's background representation function  $g$  (as discussed in Sections 4 and 5), we experimented with three variations of CFS. First, we performed the two-stage approach proposed in Section 4, where  $g$  is pretrained on background data (CFS-Pretrained). Second, we tried learning  $g$  jointly with our feature set (CFS-Joint), similar to previous work on representation learning for CA [11, 16]. Finally, in an attempt to combine the benefits of our proposed approach (i.e., improved separation of salient and irrelevant variations) with the simplicity of joint training (i.e., the ability to carry out optimization in a single step), we experimented with a modification of the joint training procedure: for target data points, we applied a stop-gradient operator to the output of the background representation function  $g$  (CFS-SG). By doing so,  $g$  is only updated for background data points and thus is not encouraged to pick up target-specific variations.
We refer the reader to Appendix D for details on model implementation and training. For all experiments we divided our data using an 80-20 train-test split, and we report the mean  $\pm$  standard error over five random seeds for each method.
6.1 Semi-synthetic data: Grassy MNIST
The Grassy MNIST dataset, originally introduced in Abid et al. [10], is constructed by superimposing pictures with the "grass" label from ImageNet [41] onto handwritten digits from the MNIST dataset [42]. For each MNIST digit, a randomly chosen grass image was cropped, resized to be 28 by 28 pixels, and scaled to have double the amplitude of the MNIST digit before being superimposed onto the digit (Figure 3a). For a background dataset we used a separate set of grass images with no digits (Figure 3b), and our goal here is to select features that best capture the target-specific digit variations.
We began by using our CFS variants with a fixed background dimension  $l = 20$  and baseline models to select  $k = 20$  features. Qualitatively, the features selected by CFS (Figure 3c; Figure S1a-b) concentrate around the center where digits are found, as desired. On the other hand, features selected by baselines (Figure 3d) are undesirably located far from the center where digits are never found (CAE), or are clustered in only a subset of locations where digit-related variations can be found (STG, DUFS), or both (LassoNet). These results illustrate the unsuitability of fully unsupervised methods for CA, as such methods may select features that contain minimal information on the target-specific salient variations. Moreover, these results illustrate that naive adaptations of supervised methods are not suitable for CA, as features that are sufficient to discriminate between target and background samples do not necessarily contain maximal information on the target-specific salient variations.
${}^{3}$  Code for reproducing our results is available in supplementary materials.
Figure 4: a, Quantitative performance comparison of CFS versus baseline supervised (S) and unsupervised (U) methods on Grassy MNIST for varying numbers of features  $k$ . b, We vary the relative contribution of grass noise to the dataset and assess each method's performance when used to select  $k = 20$  features. c, We select  $k = 20$  features on the original Grassy MNIST dataset and vary CFS's background representation dimension  $l$  to understand the impact of this hyperparameter.
Figure 5: Benchmarking CFS and baseline supervised (S) and unsupervised (U) methods on four real-world biomedical datasets.
We then evaluated the quality of features selected by each method by training random forest classifiers to predict the digit (i.e., 0-9) in each image, using only the  $k$  selected features. We report the accuracy of the classifiers on a held out test set in Figure 4a for varying values of  $k$ . We find that our pretrained and stop-gradient CFS models outperform baseline methods, with a minimal gap between them. On the other hand, CFS-Joint selected lower quality features and did not always outperform baselines, reflecting the pitfalls of joint training discussed in Sections 4-5. We next assessed how CFS's performance varies with respect to the ratio of background noise vs. salient variation by varying the scale of grass noise versus MNIST digit: when selecting  $k = 20$  features at each increment, we found that CFS continued to outperform baselines across noise scales (Figure 4b). Finally, we assessed the impact of the background representation size  $l$  for our CFS variants by selecting  $k = 20$  features on the original Grassy MNIST dataset with varying values of  $l$ , and we found that our results were largely insensitive to this hyperparameter (Figure 4c). We also found that this phenomenon was consistent across varying choices of  $k$  (Figure S3).
6.2 Real-world datasets
We next applied CFS and baseline methods to four real-world biomedical datasets: protein expression levels from trisomic and control mice that were exposed to different experimental conditions (Mice Protein) [43], single-cell gene expression measurements from mice intestinal epithelial cells exposed to different pathogens (Epithelial Cell) [44], single-cell gene expression measurements from patients with acute myeloid leukemia before and after a bone marrow transplant (AML Treatment) [45], and gene expression measurements from a high-throughput single-cell gene knockout perturbation screen (ECCITE-seq) [14]. For each target dataset, we obtained a corresponding background collected as part of the same study. We refer readers to Appendix E for further details on each dataset.
Feature quality was again assessed by training random forest models to classify target samples into known subgroups of target points, and we report classification accuracy on a held-out test set for varying numbers of features in Figure 5. Once again we find that our pretrained and stop-gradient
CFS variants outperformed baseline fully supervised and unsupervised methods, demonstrating that carefully exploiting the weak supervision available in the CA setting can lead to higher-quality features for downstream CA tasks. Furthermore, we find that CFS-Joint's performance was inconsistent across datasets, and that it sometimes underperformed baseline models. This phenomenon further illustrates the need to carefully train the background function  $g$  as discussed in Sections 4-5 to ensure the selection of features that are maximally informative for CA. Moreover, we found that these results were consistent for other choices of downstream classifier model, including XGBoost [46] and multilayer perceptrons (Figure S4). Using the mice protein dataset, we also performed further experiments assessing the importance of the background dataset for CFS's performance. In particular, we assessed the performance of our pretrained and gradient stopping CFS variants as the number of background samples varied, and we also compared against a variant of CFS where the background encoder was fixed at a random initialization (i.e., the background encoder was not updated during training). We found that CFS's performance initially improved as the number of background samples increased before saturating (Table S2), and we also found that the randomly initialized CFS variant exhibited significantly worse performance for lower numbers of features (Table S3).
Finally, to further illustrate the utility of our method for exploring real-world biomedical datasets, we performed a deeper analysis of the specific genes selected by CFS on the Epithelial Cell pathogen infection dataset, and we present these results in Appendix F. In short, we found that CFS tended to select genes that exhibited subtle cell-type-specific responses to pathogen infection. On the other hand, features selected by baseline methods were sufficient for distinguishing between control and infected cells, but did not capture these more subtle cell-type-specific effects. This result thus illustrates CFS' potential to facilitate new biological insights.
In this work we considered the problem of feature selection in the contrastive analysis (CA) setting. Our goal is to discover features that reflect salient variations enriched in a target dataset compared to a background dataset, and we tackled the problem by proposing CFS (Section 4), a method we motivated with a new information-theoretic analysis of representation learning in the CA setting (Section 5). Experiments on both simulated and real-world datasets show that features selected by CFS led to superior performance on downstream CA tasks as compared to those selected by previous feature selection methods. CFS represents a promising method to identify informative features when salient variations are subtly expressed in unlabeled contrastive datasets, and our theoretical analysis may prove useful for developing and comparing other learning methods designed for the CA setting.
This work was funded by NSF DBI-1552309 and DBI-1759487 (E.W., I.C. and S.-I.L.), NIH R35-GM-128638 and R01-NIA-AG-061132 (E.W., I.C. and S.-I.L.). E.W. was supported by the National Science Foundation Graduate Research Fellowship under Grant No. DGE-2140004.
[1] Isabelle Guyon, Jason Weston, Stephen Barnhill, and Vladimir Vapnik. Gene selection for cancer classification using support vector machines. Machine Learning, 46:389-422, 2002.
[2] Aravind Subramanian, Rajiv Narayan, Steven M Corsello, David D Peck, Ted E Natoli, Xiaodong Lu, Joshua Gould, John F Davis, Andrew A Tubelli, Jacob K Asiedu, et al. A next generation connectivity map: L1000 platform and the first 1,000,000 profiles. Cell, 171(6): 1437-1452, 2017.
[3] Muhammed Fatih Balin, Abubakar Abid, and James Zou. Concrete autoencoders: Differentiable feature selection and reconstruction. In International Conference on Machine Learning, pages 444-453. PMLR, 2019.
[4] Pabitra Mitra, CA Murthy, and Sankar K. Pal. Unsupervised feature selection using feature similarity. IEEE Transactions on Pattern Analysis and Machine Intelligence, 24(3):301-312, 2002.
[41] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248–255. IEEE, 2009.
[42] Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist, 2, 2010.
[43] Clara Higuera, Katheleen J Gardiner, and Krzysztof J Cios. Self-organizing feature maps identify proteins critical to learning in a mouse model of down syndrome. PloS One, 10(6): e0129126, 2015.
[44] Adam L Haber, Moshe Biton, Noga Rogel, Rebecca H Herbst, Karthik Shekhar, Christopher Smillie, Grace Burgin, Toni M Delorey, Michael R Howitt, Yarden Katz, et al. A single-cell survey of the small intestinal epithelium. Nature, 551(7680):333-339, 2017.
[45] Grace XY Zheng, Jessica M Terry, Phillip Belgrader, Paul Ryvkin, Zachary W Bent, Ryan Wilson, Solongo B Ziraldo, Tobias D Wheeler, Geoff P McDermott, Junjie Zhu, et al. Massively parallel digital transcriptional profiling of single cells. Nature Communications, 8(1):1-12, 2017.
[46] Tianqi Chen and Carlos Guestrin. XGBoost: A scalable tree boosting system. In Proceedings of the 22nd acm SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 785-794, 2016.
[47] Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The Concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712, 2016.
[48] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144, 2016.
[49] Emil Julius Gumbel. Statistical theory of extreme values and some practical applications: a series of lectures, volume 33. US Government Printing Office, 1954.
[50] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32:8026-8037, 2019.
[51] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[52] Md Mahiuddin Ahmed, A Ranjitha Dhanasekaran, Aaron Block, Suhong Tong, Alberto CS Costa, Melissa Stasko, and Katheleen J Gardiner. Protein dynamics associated with failed and rescued learning in the ts65dn mouse model of down syndrome. PloS One, 10(3):e0119491, 2015.
[53] Malte D, Luecken and Fabian J Theis. Current best practices in single-cell rna-seq analysis: a tutorial. Molecular systems biology, 15(6):e8746, 2019.
[54] F Alexander Wolf, Philipp Angerer, and Fabian J Theis. Scany: large-scale single-cell gene expression data analysis. Genome Biology, 19(1):1-5, 2018.
[55] Eleni P Mimitou, Anthony Cheng, Antonino Montalbano, Stephanie Hao, Marlon Stoeckius, Mateusz Legut, Timothy Roush, Alberto Herrera, Efthymia Papalexi, Zhengqing Ouyang, et al. Multiplexed detection of proteins, transcriptomes, clonotypes and crispr perturbations in single cells. Nature Methods, 16(5):409-412, 2019.
A Choice of the feature selection layer
For the experiments presented in the main text, we implemented CFS with the STG feature selection layer proposed in Yamada et al. [23]. However, CFS does not have to be implemented with STG, and can be implemented with any choice of differentiable feature selection layer. Here we describe in more detail another potential choice of feature selection layer: the concrete selector layer introduced in Balin et al. [3].
The concrete selector layer is based on concrete random variables [47, 48]. A concrete random variable can be sampled to produce a continuous approximation of a one-hot vector. To sample a  $d$ -dimensional concrete random variable, one first samples a  $d$ -dimensional vector of i.i.d. samples from a Gumbel distribution [49]  $\pmb{g}$ . Each element  $\pmb{m}$  of our concrete variable is then defined as
$$
\boldsymbol {m} _ {i} = \frac {\exp \left(\left(\log \boldsymbol {\alpha} _ {i} + \boldsymbol {g} _ {i}\right) / T\right)}{\sum_ {k = 1} ^ {d} \exp \left(\left(\log \boldsymbol {\alpha} _ {k} + \boldsymbol {g} _ {k}\right) / T\right)}
$$
Here our temperature parameter  $T \in (0, \infty)$  controls the extent to which our one-hot vector is relaxed. In the limit  $T \to 0$ , our concrete random variable outputs one-hot vectors with  $m_i = 1$  with probability  $\alpha_i / \sum_{k=1}^d \alpha_k$ . Concrete random variables are differentiable with respect to their parameters  $\alpha$  via the reparameterization trick [39].
The concrete selector layer then selects features in the following manner. For each of the  $k$  nodes in the selector layer, we sample a  $d$ -dimensional concrete random variable  $\pmb{m}^{(i)}$ ,  $i \in \{1, \dots, k\}$ . The  $i$ th node in the selector layer  $u^{(i)}$  outputs  $\pmb{x} \cdot \pmb{m}^{(i)}$ . As  $T \to 0$  this inner product simply becomes equal to one of the input features. After the network is trained, the concrete selector layer is replaced by an arg max layer for which the output of the  $i$ th neuron is  $\pmb{x}_{\arg \max_j \pmb{\alpha}_j}$ .
Before training, the selector's parameters  $\alpha_{j}$  are initialized as small positive values to encourage the layer to explore different combinations of input features. As the layer is trained, the values of  $\alpha_{j}$  become more sparse as the layer becomes more confident in particular choices of features.
Similar to [3], we can train concrete selector layers with CFS using a simple annealing schedule for the temperature parameter  $T$ . That is, we begin training with a high temperature  $T_{0}$  and gradually decay the temperature until we reach a final temperature  $T_{B}$  at each epoch according to a first-order exponential decay  $T(b) = T_{0}(T_{B} / T_{0})^{b / B}$  where  $T(b)$  is the temperature at epoch  $b$  and  $B$  is the total number of epochs.
In Figure S2 we present the results of applying CFS with the concrete selector layer to select features on the Grassy MNIST dataset. As with our original CFS implementation with the STG layer, we found that the pretrained and gradient-stopping CFS variants implemented with the concrete layer selected features around the center of images as desired.
In this section, we re-state and then prove our theoretical results from the main text.