The field of data science has seen the growing importance of dimension reduction (DR) techniques as a preliminary step in processing large-scale biological datasets, such as single-cell RNA sequencing data. These techniques aid in tasks like data visualization, structural pattern discovery, and subsequent biological analyses. Within this broader context, Supervised Dimension Reduction (SDR, also known as sufficient dimension reduction) methodologies have gained significant attention [16, 37]. In the study by [61], a comparison is drawn between SDR techniques and unsupervised counterparts like Principal Component Analysis (PCA), Spherelets [40], and Spherical Rotation Component Analysis [46], emphasizing the increasing prominence and utility of SDR in contemporary data science.

Given paired observations ( x , y ) ∈ R p × R, where x consists of p covariates, and y is the corresponding response or output variable, the common assumption in SDR is that (1.1) y = φ ( V ⊤ x , ϵ ) , for some function φ , where V ∈ R p × d with d ≪ p is the projection matrix from a high-dimensional to a low-dimensional space, ϵ is the measurement error independent of x, and φ is an arbitrary unknown function. For example, in a single-cell RNA sequencing dataset, x could be the expression of genes for a cell and y could be the cell type.

Under assumption (1.1), although the low-dimensional representation V ⊤ x is determined by a linear transformation, the function φ is an arbitrary unknown function. In this paper, we stick to the assumption in (1.1) to focus on linear DR methods for two reasons. First, linear methods are computationally more efficient, particularly for large p and large n. Second, linear methods are more interpretable, which is an essential characteristic in scientific applications. For instance, in the example above, each column of V ⊤ x is often considered as a genetic pathway [4]. Although our proposed method can be extended to nonlinear cases by the kernel method, we will leave this for future work.

Sliced Inverse Regression (SIR, [41]) is a well-established technique for supervised dimension reduction that is widely applicable in multiple scenarios due to its roots in regression analysis. It has been shown to have strong consistency results in both fixed dimensional [32] and high-dimensional [45] settings. The goal of SIR is to capture the most relevant low-dimensional linear subspace without any parametric or nonparametric model-fitting process for φ.

Moreover, SIR offers a geometric interpretation by conditioning on the sufficient statistics of the input distribution [41, 18]. SIR incorporates the idea of linear dimension reduction with statistical sufficiency. In SIR, given a pair of features x ∈ R p and univariate response y ∈ R, the goal is to find a matrix V ∈ R p × d , d < p such that y is conditionally independent of x given V ⊤ x. Although the matrix V is not identifiable, the column space of V, denoted C ( V ), is identifiable.

Motivated by emerging modern high-dimensional [25, 44, 64] and biological datasets [28, 42], SIR evolved and admitted several generalizations, including localized SIR [68], kernel SIR [67], SIR with regularization [42], SIR for longitudinal data [34, 43], metric response values [60], and online SIR [10].

In this article, we focus on a specific type of high-dimensional biological data, where the dataset consists of two groups — a foreground group, also known as treatment group or case group, and a background group, also known as control group. The goal is to identify the low-dimensional structure, variation, and information unique to the foreground group for downstream analysis. This situation arises naturally in many scientific experiments with two subpopulations. For example, in Electronic Health Record (EHR) data, the foreground data could be health-related covariates from patients who received certain medical treatment, while the background data could be measurements from healthy patients who did not receive any treatment. In this case, the goal is to identify a unique structure in patients who received the treatment that can predict future outcomes. In a genomics context, the foreground data could be gene expression measurements from patients with a disease, and the background data could be measurements from healthy people. In this case, the goal is to predict a certain phenotype for the diseased patient for disease analysis and future therapy.

Previous contrastive models, such as the contrastive latent variable model (CLVM, [73]), contrastive principal component analysis (CPCA, [1]), probabilistic contrastive principal component analysis (PCPCA, [38]), and the contrastive Poisson latent variable model (CPLVM, [35]), have shown that using the case-control structure between foreground and background groups can greatly improve the effectiveness of dimension reduction over standard DR methods such as PCA and its variants. However, to the best of our knowledge, none of these unsupervised contrastive dimension reduction methods is directly applicable to the SDR setting.

In this work, we move from these unsupervised contrastive dimension reduction methods to a supervised contrastive dimension reduction setting. By combining the idea of contrastive loss function and the sufficient dimension reduction considered in the SIR model, we propose the Contrastive Inverse Regression (CIR) model, which exactly recovers SIR when a certain parameter is zero. The CIR model sheds light on how to explore and exploit the contrastive structures in supervised dimension reduction.

Table 1 lists several popular DR methods and their properties. The table categorizes these methods as “contrastive” and “supervised”, based on whether they are designed for case-control data and able to identify low-dimensional structure unique to the case group, and if they take the response variable y into consideration and use V ⊤ x to predict y. For example, PCA, the most well known DR method, is neither contrastive nor supervised. Similarly, canonical correlation analysis (CCA, [31]) does not utilize y or the unique information of one group, which makes it neither contrastive nor supervised. Methods such as CLVM, CPCA, PCPCA, and CPLVM are contrastive but not supervised. On the other hand, classical supervised DR methods including SIR [41], linear discriminant analysis (LDA, [26]), and the least absolute shrinkage and selection operator (LASSO, [58]) are supervised but not contrastive. Our proposed method, CIR, combines both contrastive and supervised features by utilizing both the response y and the case-control structure.

It is important to note that the assumption (1.1) does not limit the response variable y to be continuous or categorical, and thus we do not distinguish between regression and classification. However, some methods listed in Table 1 are specifically designed for either continuous y (regression, LASSO) or categorical y (classification, LDA). CIR handles both scenarios with the only difference being in the choice of slices, as explained in Section 2. Furthermore, not all existing DR methods are included in this table. For example, the recently proposed linear optimal low-rank projection (LOL, [61]) is designed for the classification setting and requires the number of classes to be smaller than the reduced dimension d. This can be restrictive, for example, when applied to a single-cell RNA sequencing dataset, where d is required to be greater than the number of cell types. In contrast, CIR does not require such data-dependent constraints on the reduced dimension d. Similarly, data visualization algorithms that require d = 2 such as the t-distributed stochastic neighbor embedding (tSNE, [59]) and uniform manifold approximation and projection (UMAP, [5]) are not listed in the table.

We now present our proposed methodology, including an algorithm for solving the associated nonconvex optimization problem on the Stiefel manifold. We also provide analysis of the convergence of the algorithm, and conduct extensive experiments to demonstrate its superior performance on high-dimensional biomedical datasets when compared to existing DR methods. All proofs are provided in the appendix, and additional experimental details are in the supplement material.

To maintain consistency, we will use the terms “foreground group” and “background group” instead of “case-control” or “treatment-control” in the remaining sections. First, we briefly review SIR as our motivation.

Recall that under the assumption in Equation (1.1), the centered inverse regression curve, E [ x | y ] − E [ x ], lies exactly in the linear space spanned by columns of Σ x x V, denoted by C ( Σ x x V ), where Σ x x is the covariance matrix of x. This linear subspace is called the effective dimension reduced (e.d.r.) space [41]. As a result, the objective of SIR is to minimize the expected squared distance between E [ x | y ] and C ( Σ x x V ): (2.1) min V ∈ St ( p , d ) E y [ d 2 ( E [ x | y ] − E [ x ] , C ( Σ x x V ) ] where d is the Euclidean distance.

In the contrastive setting, denote foreground data by ( x , y ) ∈ R p × R and background data by ( x ˜ , y ˜ ) ∈ R p × R. For convenience, we assume that x and x ˜ are centered at the origin so that E [ x ] = E [ x ˜ ] = 0.

Our goal is to find a low-dimensional representation of x, denoted by V ⊤ x, such that y is determined by V ⊤ x while y ˜ is not determined by V ⊤ x ˜. The goal of CIR is to find a low-dimensional subspace represented by V such that (2.2) y = φ ( V ⊤ x , ϵ ) , for some function φ , y ˜ ≠ ψ ( V ⊤ x ˜ , ϵ ˜ ) , for any function ψ . That is, the column space of V captures the low-dimensional information unique to the foreground group so that we can use V ⊤ x to predict y through φ, but cannot use V ⊤ x ˜ to predict y ˜ through any function ψ. Instead of optimizing a single loss similar to SIR, CIR aims at optimizing the subspace C ( Σ x x V ) to minimize the “contrastive loss function”: (2.3) f ( V ) : = E y [ d 2 ( E [ x | y ] , C ( Σ x x V ) ] − α E y ˜ [ d 2 ( E [ x ˜ | y ˜ ] , C ( Σ x ˜ x ˜ V ) ] , where α ≥ 0, Σ x x = Cov ( X ) and Σ x ˜ x ˜ = Cov ( X ˜ ) ∈ R p × p , and d is the Euclidean distance. Define the following notation: v y = E [ x | y ] , v y ˜ = E [ x ˜ | y ˜ ] ∈ R p Σ x = Cov ( v y ) , Σ x ˜ = Cov ( v y ˜ ) ∈ R p × p v y (and v y ˜ ) are called the centered inverse regression curves [41, 60]. The resulting loss function f balances the effectiveness of dimension reduction between the foreground and background groups. We can adjust the hyperparameter α to express our belief in the importance of the background group. Note that the parameter α appears naturally in other contrastive DR methods, including CPCA and PCPCA.

Note that V is not identifiable, and is identifiable only up to a d-dimensional rotation. However, the contrastive loss f, determined by V V ⊤ , the projection matrix to the column space of V, is invariant under such rotations. This nonidentifiability issue is common in other DR methods, including PCA, CPCA, SIR, etc, where the convention is to refer to the column space of V. Therefore, this nonidentifiability is consistent with standard practices and does not impact the validity of our results.

Observe that in the case where α = 0, CIR reduces to SIR. In this case, the problem can be reparameterized by V ∗ = B 1 / 2 V so that the columns are orthogonal, which reduces the loss function to a quadratic form, yielding a closed-form solution (as a generalized eigenproblem). In the case where α > 0, however, we cannot perform the same trick for both B and B ˜, so we must resort to numerical approximations. We adopt gradient-based optimization algorithms on St ( p , d ), which are based on the gradient of f given by the following lemma.

Note that the gradient grad f is different from the standard gradient in Euclidean space, denoted by D f = ∂ f ∂ V . The difference is that grad f lies in the tangent space of St ( p , d ) at V, while the Euclidean version may escape from the tangent space.

Let G ( V ) = V ⊤ A V and G ˜ ( V ) = V ⊤ A ˜ V. Note, then, that the local optimality condition is equivalent to (2.4) A V E ( V ) − α A ˜ V E ˜ ( V ) = B V E ( V ) G ( V ) E ( V ) − α B ˜ V E ˜ ( V ) G ˜ ( V ) E ˜ ( V ) .

In Appendix G, we discuss how Equation (2.4) may lead to a gradient-free algorithm that involves solving an asymmetric algebraic Ricatti equation.

So far, we have discussed the population version, which relies on the distributions of x, x ˜, y, and y ˜ that are unknown in practice. In real applications, we observe finite samples ( x i , y i ) i = 1 n as foreground data and ( x ˜ j , y ˜ j ) j = 1 m as background data. We denote X ∈ R n × p , X ˜ ∈ R m × p where each row represents a sample; similarly, each entry of Y ∈ R n and Y ˜ ∈ R m represents a response value. In this case, we can replace the expectation by the sample mean to get estimates of Σ x , Σ x ˜ , Σ x x , and Σ x ˜ x ˜ and have the corresponding plug-in estimates for A, B, A ˜, and B ˜. Throughout this paper, we assume p < n so that all covariance matrices are nonsingular. The extension to p > n is discussed in Section 5. After computing these matrices, the problem is reduced to a manifold optimization problem [2]. The estimates of Σ x and Σ x ˜ deserve further discussion. As shown by [41] and [10] among others, for continuous response y, the observed support of response y can be discretized into slices I h = ( q h − 1 , q h ], for h = 1 , … , H. An estimate of Σ x is given by ∑ h = 1 H m h m h ⊤ where m h = 1 n p h ∑ y i ∈ I h x i with p h = 1 n ∑ i = 1 n I ( y i ∈ I h ). On the other hand, if y and y ˜ are categorical, the slices are naturally chosen as all possible values of y and y ˜. Combining these pieces, we present our empirical version of the CIR algorithm in Algorithm 1.

It is worth noting that our optimization of f as a function of V ∈ St ( p , d ) cannot be considered as an optimization problem in R p × d with orthogonality constraints V ⊤ V = I d  [23, 7]. Because the term ( V ⊤ B V ) − 1 in f is not well defined unless V is full rank, our loss function f cannot be extended to the full Euclidean space R p × d . We consider it as an optimization problem intrinsically defined on St ( p , d ) as laid out by [2]. This key property excludes some commonly used optimizers on manifolds, and we will discuss more details in the next section.

In this section, we discuss two concrete optimization algorithms for the last step in Algorithm 1 to find V ∗ and show their convergence. The optimization problem outlined in Equation (2.3) does not follow the classic Li-Duan theorem for regression-based dimension reduction due to its nonconvex nature (see, e.g., [17, Prop. 8.1]). The convergence of the algorithm is discussed in detail below.

The first algorithm we consider is the scaled gradient projection method (SGPM) specifically designed for optimization on the Stiefel manifold [50]. We first define an analog to Lagrangian multiplier L ( V , Λ ) : = f ( V ) − 1 2 ⟨ Λ , V ⊤ V − I d ⟩, then the SGPM algorithm is summarized in Algorithm 2, where π ( X ) = arg min Q ∈ St ( p , d ) | | X − Q | | F is the orthogonal projection to the Stiefel manifold. Note for Algorithm 2 that the update for μ and C k + 1 is intricate; see [50] for more details.

To study the convergence of Algorithm 2, we need to study the Karush-Kuhn-Tucker (KKT) conditions for CIR:

Now we can state the convergence theorem of Algorithm 2:

Although Algorithm 2 is guaranteed to converge, there are two drawbacks. First, the accumulation point V ∗ is only guaranteed to satisfy the KKT conditions, but not necessarily a critical point. Second, we do not know how fast V k will converge to V ∗ . Next, we introduce an accelerated line search (ALS) algorithm as an alternative to SGPM that converges to a critical point with a known convergence rate. ALS is summarized by Algorithm 3, where t k A is the step size, called the Armijo step size for given α ‾, β, σ, η k and R is a retraction to St ( p , d ), see [2] for more details.

Algorithm 3 can be shown to have linear convergence to critical points if the hyperparameters are chosen properly. For other choices of η k , see Appendix E for more details. The following adaptation of Theorem 4.5.6 in [2] indicates linear convergence to stationary points.

The difference between Algorithm 2 and 3 deserves further comment. While we empirically observe that Algorithm 2 often converges faster, Algorithm 3 has theoretical properties which allow for a proof of linear convergence in terms of an upper bound on the number of iterations. In practice, for smaller datasets we suggest running Algorithm 3, while for larger datasets we recommend using Algorithm 2 for efficiency.

The computational complexity of CIR for both the SGPM-based optimization and the ALS-based optimization is compared to various competitors in the table below. Here, we assume that 1 ≤ d < p < m , n, where the background and foreground data have m and n samples, respectively. Specifically, we assume that ϵ denotes the stopping error such that f ( V k ) − f ( V ∗ ) ≤ ϵ. It is noteworthy that a p-dimensional singular value decomposition can be achieved within O ( p 3 ). We present the comparison in the table below.

When applying CIR, several hyperparameters must be tuned, such as the weight α, the reduced dimension d, and the slices I h and I h ˜ for estimation of Σ x and Σ x ˜ . In some cases, it may also be necessary to determine the definition of the foreground and background groups and to assign background labels Y ˜.

The value of α ≥ 0 can be determined by cross-validation. In particular, we suggest the following 2-step method. First, identify the rough range of α at the logarithmic scale. Because CIR coincides with SIR when α = 0, running SIR initially can provide insights: better performance of SIR suggests a smaller α, and vice versa. Once a rough range is identified, standard cross-validation can be used within this range. Our numerical experiments show that CIR is robust to the choice of α; that is, the performance of the method changes continuously with α. Tables supporting this observation are provided in the supplement material.

Additionally, the choice of reduced dimension d may depend on the goal of the investigator. If visualization is considered important, d = 2 is appropriate. If the goal is prediction, the elbow point of the d versus prediction error plot may suggest an optimal d. However, as with other DR methods, determining the optimal value of d is still a topic of ongoing research [12, 13].

The definition of foreground data X and Y should be the data and the target variable of interest, while the choices of background data X ˜ and Y ˜ may not be as straightforward. These data are intended to represent “noise” that is “subtracted” from the foreground data. For example, in the biomedical context, if the population of interest is a group of sick patients, the background dataset may include observations of healthy individuals. In other contexts, however, it may be appropriate to use X ˜ = X. In this case, the choice of background label Y ˜ may be unclear. If another outcome variable was collected, it could be used as Y ˜; otherwise, randomly selected values in the support of Y could be used to represent “noise”.

The estimates for Σ x and Σ x ˜ are partly determined by whether y and y ˜ are categorical or continuous. If these variables are categorical, then each value of y (or y ˜) can be considered as a separate slice, resulting in | supp ( Y ) | (or | supp ( Y ˜ ) |) total slices. On the other hand, if these variables are continuous, slices can be taken to represent an equally spaced partition of the range of Y (or Y ˜), with the number of slices being tunable hyperparameters.

In this work, we propose the CIR model and the associated optimization algorithm for supervised dimension reduction for datasets that are split into foreground and background groups. We provide a theoretical guarantee of the convergence of the CIR algorithm under mild conditions. We have shown that our CIR model outperforms competitors in multiple biomedical datasets, including mouse protein expression data, single-cell RNA sequencing data, and plasma retinol data. However, there are several important future directions that remain unaddressed in this paper, as outlined below.

Multi-Treatment  It is natural to consider how our model can be extended to studies with multiple treatments. For example, in medical treatment, there might be more than one treatment for patients with certain disease. In [61], it has been shown that the number of treatment groups puts a hard constraint on the target dimension. It is interesting to generalize from a single-treatment structure to a multi-treatment structure (e.g., [47]), where the loss function needs more sophisticated design.

Another direction in multi-group scenario is to combine multiple CIR models trained on different pairs of bi-groups. As pointed out by [65], the generalization error in the contrastive regression model stacking needs to be controlled, and one possible way is to follow divergence mixing as proposed by [29], with a careful normalization. The major difficulty in training such stacking model is how to devise a sequential optimization for model training.

Consistency and Sufficient Dimension Reduction  The consistency of the proposed CIR model remains open. Theorems 2 and 3 ensure that the resulting solution must be a stationary point, but we did not discuss whether these stationary points are consistent estimates. The consistency of the estimates is also affected by the choice of α, because α = 0 will render this CIR into a classical SIR for the foreground group. This consistency problem also has practical importance, as it explicitly expresses the trade-off between the expressive contrastiveness and the emphasis on the effective lower-dimensional structures. The group information and statistical sufficiency compete against each other, as we observed in the experiments, thus a range of α that balances between these two factors are of interest and might be answered by the consistency result.

Furthermore, SIR has the drawback of missing the totality central subspace when the symmetry assumption in x is lost [37]. [18] proposed the sliced average variance estimator (SAVE) estimator for addressing this problem, which raises the natural question of how to generalize this high-moment SDR method to the contrastive setting.

Loss Function  Our loss function (2.3) is nonstandard, which raises many questions. For example, the relation between number of local minima and A, B, A ˜, B ˜, α remains open. Moreover, although f cannot be continuously extended to the full Euclidean space R p × d , if we restrict the domain to be a submanifold of St ( p , d ), f might be extended to the convex hull of the submanifold. This extension will enable us to apply some other efficient optimization algorithms with strong theoretical guarantee [7]. Furthermore, Appendix G raises the question of the validity of a fixed-point algorithm based on Ricatti equations that may lead to a more efficient algorithm to minimize f without involving the gradient.

Scalability  The method we presented in this paper does not handle high-dimensional data in the sense of p > n , m, because matrices B and B ˜ are singular in this situation. A possible extension of CIR to p > n is to use the same technique as sparse PCA [72], which introduces a penalty term to enforce sparsity.