Disentangling Interpretable Factors with Supervised Independent Subspace Principal Component Analysis

Response to Reviewer debZ

Thank you for your positive feedback and insightful comments! We appreciate your thorough review and would very much like to address any remaining concern.

W1. Impact of $\lambda$ selection and an auto-selection pipeline

Thank you for the suggestion. We conducted additional experiments by varying $\lambda$ for the scRNA-seq dataset with values [0, 0.1, 0.3, 1, 3, 10, 30, 100]. On top of that, we implemented a preliminary pipeline for automatic $\lambda$ selection that aggregate resulting subspaces using clustering and outputs aggregated results at different resolution, as implemented in Contrastive PCA. However, due to time constraints, we have yet to thoroughly test some design choices, such as the representation distance, which appears to be important in subspace clustering. To speed up training, we also introduced a mini-batch approach of Alg 1 (Appendix B.2) since the learnable projection matrix U is of size (n_feature, n_dim), independent of n_sample. We plan to implement additional features such as parallel training and better initialization search in the future.

Despite these limitations, we confirmed the following results:

Impact on subspace similarity

Subspaces learned with similar $\lambda$ values are generally more similar, as measured by Grassmann distance between representations, indicating gradual and continuous changes as a function of $\lambda$. The infection subspaces for $\lambda$ = [0.1, 0.3, 1] cluster together, as do the time subspaces for $\lambda$ = [0, 0.1], while the other subspaces appear more detached.

Impact on gene selection and GO analysis

We quantified overlaps of the top 100 genes list for the infection subspace (ranked by absolute PC1 weights) across different $\lambda$ values. Results for $\lambda$ = [0, 0.1, 0.3, 3, 10] are highly consistent, sharing more than 90 of the top 100 genes. The result also suggests that sPCA ($\lambda$ = 0) is sufficient for the GO analysis in Figure 5 ($\lambda$ =10). We examined sPCA representations and confirmed that the confounding effects from the temporal subspace (as visualized in Figure 9b and the new Fig S1c in M2) are mainly concentrated in PC2, which is not used here for feature selection. Notably, for reasons still under investigation, $\lambda$ = 1 is an outlier, sharing only 68 of the top 100 genes with sPCA. A possible explanation is that Alg 1 (Appendix B.2) on large datasets may experience numerical issues with PyTorch’s SVD solver.

We will add visualizations of these results in the final version.

W2. Connection with self-supervised learning

We apologize for any ambiguity. Our reference of "self-supervised learning" is the connection to the auto-encoder (Appendix B.1), where target variables - features from the data - are used to disentangle the data itself. We chose the term mainly to reflect the difference with supervised learning. Here our goal isn't perfect target prediction, as labels are already known. Instead, we aim to better explain the data by reweighting the self-reconstruction loss with supervision guidance (as reflected in line 401). We recognize the potential confusion and will clarify the term in the final version. The connection to regularized linear regression is discussed in Appendix B.3.

Minor points

Thank you for pointing out the problems.

• We will fix citation issues, and will increase the DPI of all figures to improve resolution.
• We will provide more details on the preprocessing of the single-cell infection data. The data, preprocessed by the original authors, was subject to library size normalization followed by log1p transformation. A background correction step was also applied before normalization, where the mean expression in empty wells was subtracted. This makes the raw counts technically not integers, but the negative binomial likelihood in hsVAE-sc (see M2) seems to work just as well.

Response to Reviewer QSK9

Thank you for your thoughtful review of our work. We address your main points below:

Weaknesses

W1: Theoretical analysis and subspace recovery

We acknowledge the limitations in our theoretical analysis. However, sisPCA makes minimal assumptions about the data, and its objective function provides straightforward insights into the model's behavior and optimization landscape (Conjecture 3.1 and Appendix D).

Regarding the simulation example

1. We confirmed that PCA-based models yield similar results with Gaussian-drawn mixing matrices compared to the uniform case. This is because the mixing matrix was normalized to ensure equal contribution from each subspace.
2. As discussed in Appendix D (lines 449-453), subspace scale and supervision strength are the major factors influencing the optimization landscape. For example, sisPCA will prioritize S2’s structure and ignore other subspaces if we scale up the S2 supervision (X, Y) large enough.

General subspace recoverability depends on supervision quality

1. In the most intertwined case of identical supervision for two subspaces, only one can be recovered and the other will collapse to rank zero (Figure 6, Appendix D). Here, linearity ensures that the same information is not split between subspaces - it's an all-or-nothing scenario.
2. When both supervised and unsupervised subspaces are presented, their dynamics are complicated by the dual role of HSIC supervision indicated in Eq.4, Appendix B.3. In particular, the supervised subspace will also try to expand upon the direction that maximize the (unsupervised) variance, potentially leading to identifiability or multi-optima issues (Section 3.3, lines 190-193, Figure 7). Though we reason that the same information will again concentrate in one subspace due to linearity.

W2: Comparison with state-of-the-art baselines

We appreciate your suggestion to include more state-of-the-art methods in our comparison. In response:

1. We've added comparisons with non-linear VAE-based models (M2), focusing on the real single-cell data where the practical need is to learn genes responsible for malaria infection defense (a feature selection task).
2. Our analysis shows that while VAE-based models can extract independent subspaces, they lack a straightforward method for identifying genes upon which each subspace was constructed (M2.2).
3. Regarding representation quality, we've confirmed that linear models like sPCA serve as strong baselines, both qualitatively and quantitatively (see Fig S1, Table S2 in M2.3). Specifically, quality of the linear infection subspace often surpasses its non-linear counterpart.

General questions

Q1: Effect of non-linear kernels for target variables

Yes, the choice of target kernel does influence sisPCA's results, primarily affecting the rank of the learned subspace. This is because sPCA's subspace dimensions are determined by supervision kernel rank (Appendix B.3, line 406-408), and sisPCA shows similar behavior. Specifically, for a linear kernel on K continuous features, the effective dimension of the corresponding sisPCA subspace is the rank of the sum of a rank-K kernel and a rank-D disentanglement kernel, where D depends on other subspace dimensions (line 146-147). A non-linear target kernel can effectively expand the learned subspace while maintaining interpretability, albeit losing the connection with regularized linear regression (Appendix B.3).

Q2: Clarification on Fig. 2 and maximization objectives

We apologize for the confusion. To clarify:

• A and B represent different target variables used for supervision.
• The maximization arrow within each subspace represents a PCA-like objective of variance maximization, which is equivalent to maximizing the HSIC towards an identity kernel (Appendix B.1, line 397-398).
• In supervised PCAs, the variance maximization objective comes from the HSIC to the target (first component of Eq.2, see Eq.4 in Appendix B.3). Upon decomposition, the arrows between the subspace and A/B indeed represents minimizing the prediction error (first component of Eq.4), and the arrows within each subspace represents maximizing the variance (second component of Eq.4).

We will revise the figure to make notation consistent.

Q3: Definition of balanced/unbalanced supervision

Balanced/unbalanced supervision refers to the relative ratio of scale/norm between target kernels (discussed in Appendix D, line 449-473). It can alternatively be defined as the ratio of loss gradient with respect to each supervision kernel. Intuitively speaking, unbalanced scenarios occur when one kernel is scaled up, favoring the corresponding subspace and creating a single global optimum (Figure 6 in Appendix D).

Responses to Reviewer UeQu

Thank you for your constructive feedback on the strengths and weaknesses of our work. We address your main concerns below:

Weaknesses

W1: Limited novelty

While our method extends PCA, it offers significant contributions by uniquely combining supervision and disentanglement in a linear, interpretable framework valuable for feature selection. This addresses an important gap in current dimensionality reduction techniques, especially in an era where linearity and interpretability are often overlooked (See M1).

W2: Comparison with SOTA methods

We've added comparisons with non-linear VAE-based models (See M2), focusing on:

1. Representation quality: Low-dimensional representations capturing specific data aspects.
2. Interpretability: Understanding feature contributions to different data aspects.

While the interpretability criterion clearly flavors linear models, our results show that linear models like sPCA are indeed strong baselines too in representation quality (See M2.3, Fig S1, Table S2). We also confirmed similar insights on supervision and disentanglement regularization effects across both linear and non-linear models.

W3: Inadequate discussion of limitations

We apologize for the poor presentation of limitations, which are somewhat scattered throughout the paper. We will concentrate them and improve the discussion section in the final version. Key limitations and their potential impacts include:

1. The limitation of linear-kernel HSIC on theoretical independence guarantee (Section 3.1, line 159 - 165). Though in practice this is less concerning, since minimizing HSIC-linear will also reduce HSIC-Gaussian (Section 4.2.3, Table 1).
2. Subspace interpretability issue (Section 3.3, line 189 – 193, and Appendix D on the resulting optimization landscape). This would make it especially challenging when the model needs to learn both supervised and unsupervised subspaces. Figure 7 is an example where unsupervised subspace failed to capture the true dimension.
3. Challenges in balancing the disentanglement regularization strength \lambda (Section 5). This is a general challenge to all supervised disentangled models.
4. Potential performance loss due to the lack of non-linear feature interactions (Section 3.1, line 166-169; Section 5). This could lead to underperformance on datasets with complicated patterns.

General questions

Q1: How to set model hyperparameters, in particular the dimension of each subspace.

We again apologize for the lack of clarity. The main hyperparameters are subspace dimensions and disentanglement penalty scale λ (See Table S1, M2.1). The latter was discussed as a limitation in the Discussion and above in W3. For sPCA and sisPCA, subspace dimension generally doesn't affect performance beyond numerical issues from SVD solvers.

• sPCA's subspace dimensions are determined by supervision kernel rank (Appendix B.3, line 406-408). That is, the effective dimension is always 1 for a 1-D continuous variable with linear kernel (age in Section 4.2) or K for a categorical variable of K groups with delta kernel (cancer type in Section 4.2) regardless of the specified dimension. The extra axes beyond the effective dimension (sPCA-age PC2-10) will collapse to zero since the eigenvalue is zero.
• The effective dim of sisPCA subspaces is determined by both the supervision kernel rank as well as other subspaces (Remark 3.2, line 144-147). In practice, sisPCA shows similar behavior as sPCA, with extra dimensions typically collapsing to zero. For example, the sisPCA-infection subspace in Section 4.3 has approximately two effective dimensions (Figure 4, 9 and Figure S1).
• In contrast, VAE models are more sensitive to dimension changes and have more hyperparameters to tune (M2.1), which is generally beyond the scope of our work. Due to time constraints, we mostly follow the SCVI’s default in designing VAE models for benchmark.

Q2: Better description of experiments

Simulated data Figure 3

As mentioned in Q1, the dimensions of supervised subspaces are mostly predetermined by their target kernel ranks. Here dim(S1) and dim(S2) are 2 regardless of hyperparameter choice. The dimension of the unsupervised space S3 is unknown, and we aim to recover the ground truth donut structure (dim = 2) as did in Fig 3c. Figure 7 represents a failed example where only one dimension is recovered.

Subspace dimensions in Section 4.2 and 4.3

Addressed in Q1. Dimensions in PCA-based models are generally ranked by importance and can be selected using methods like "the elbow curve" (Discussion, line 311-313). In M2 subspace dimensions of VAEs models are set to 10 to ensure fair comparison.

Applications to other types of real-world datasets

We note that similar to PCA and its extensions [1], our model is also general-purpose and can assist with data exploration tasks (as pointed out by Reviewer debZ). The only domain-specific model is the hsVAE-sc tailored for single-cell data with a count-based likelihood.

[1] Abid et al., "Contrastive principal component analysis"

Q3: The use of HSIC as a quality metric in Section 4.2

The HSIC metrics in Table 1 are to support the claim that minimizing the HSIC-linear can also reduce the computationally more expensive HSIC-Gaussian (not explicitly optimized for in the loss) and thus encourage independence in the representations (line 259-261).

We appreciate your feedback and apologize for our previous insufficient response. Due to space constraints, we couldn't include additional figures and results initially. To further illustrate sisPCA's versatility as a plug-in replacement of PCA for disentangling quantities, we provide below a new application, which will be included in the appendix of the final paper, and as a tutorial upon package release.

Problem and Dataset

We used the Kaggle Breast Cancer Wisconsin Data Set (569 samples, 30 summary features, uciml/breast-cancer-wisconsin-data, CC BY-NC-SA 4.0 license). The 30 real-valued features are computed from imaging data of breast mass, which include the mean, standard error and the extremum of quantities like cell nuclear radius, texture, perimeter etc.

Here, our goals are to:

1. Learn a representation for predicting sample disease status (Malignant or Benign, not used during training).
2. Understand how original features contribute to data variability, the learned representation, and diagnosis potential.

Experiments and Results

We focus our comparison on three linear models: PCA, sPCA, and sisPCA, using zero-centered and variance-standardized data as inputs. The diagnosis label used to measure subspace quality remains invisible to all models. Below we summarize the quantitative results, and will include subspace visualization in the final paper.

PCA

We projected all features (dim = 30) onto one PCA subspace (dim = 6, determined by the elbow rule), explaining 61.6% of total variance. Malignant and benign samples appear well separated in PC1 and PC2. 'symmetry_mean' (loading = -0.223) and 'radius_mean' (loading = -0.219) are the top 2 features that negatively contribute to PC1. That is, the higher the two feature scores, the lower PC1 score and the greater the possibility of the sample being malignant.

sPCA

From the PC1 loadings, we sought to construct two subspaces to separately reflect nuclei size (‘radius_mean’ ) and shape ('symmetry_mean'). We set $Y_{radius}$ ('radius_mean', 'radius_sd') of $569 \times 2$ as the target variable for the radius subspace, and $Y_{symmetry}$ ('symmetry_mean', 'symmetry_sd') of $569 \times 2$ as the target for the symmetry subspace. The remaining 26 features were projected onto the two subspaces using sPCA (dim = 3, effective dim = 2). Both subspaces better explained diagnosis status than PCA (Table S3) but remained highly entangled. Specifically, the PC2 loadings of the two spaces have a Pearson correlation of 0.716, and 'perimeter_worst', which is highly correlated to 'radius_mean' (corr = 0.965), also strongly contributes to the PC2 of the symmetry subspace (loading = 0.238).

sisPCA

We next applied sisPCA ($\lambda$ = 10) to further disentangle the two subspaces, increasing their separation (Grassmann distance increased from 1.593 to 2.058, PC2 loadings correlation decreased to 0.257). Here ‘perimeter’ features no longer contribute to the symmetry subspace. As a result, the radius subspace remained predictive of diagnosis, while the symmetry subspace became less relevant (Table S3). We confirmed the finding by directly measuring the predictive potential of the target variables $Y_{radius}$ and $Y_{symmetry}$ (Silhouette score = 0.457 and 0.092, respectively).

Table S3: Predictability of diagnosis status, measured by Silhouette score

Interpretation

1. Our results suggest that nuclear size (radius subspace) is more informative for breast cancer diagnosis than nuclear shape (symmetry subspace), aligning with clinical observations [1].
2. Unsupervised PCA and sPCA captured both radius and symmetry aspects as they contribute most to data variability. Without disentanglement, the two aspects remain intertwined, potentially leading to incorrect conclusions about symmetry's predictive power for diagnosis.
3. sisPCA successfully separated the two representations, revealing their distinct relationships to diagnosis. The results are interpretable: the radius subspace is constructed using features like 'area' and 'perimeter', while the symmetry subspace uses features like 'compactness' and 'smoothness'.

This example demonstrates sisPCA's ability to disentangle different aspects of data variation and uncover underlying relationships. Importantly, sisPCA improves upon the diagnosis potential of PCA's representation, even through the diagnosis labels were never used during training.

[1] Kashyap, Anamika, et al. "Role of nuclear morphometry in breast cancer and its correlation with cytomorphological grading of breast cancer: A study of 64 cases." Journal of cytology 35.1 (2018): 41-45.

Response to Reviewer PWKM

Thank you for your summary. We appreciate your evaluation that our technical contribution is novel enough to warrant presentation at NeurIPS.