Invariant Feature Subspace Recovery for Multi-Class Classification
Abstract
Domain generalization aims to learn a model over multiple training environments to generalize to unseen environments.
Recently, Wang et al.
[2022] proposed
Invariant-feature Subspace Recovery (ISR), a domain generalization algorithm that uses the means of class-conditional data distributions to provably identify the invariant-feature subspace under a given causal model.
However, due to the specific assumptions of the causal model, the original ISR algorithm is conditioned on a single class only, without utilizing information from the rest of the classes.
In this work, we consider the setting of multi-class classification under a more general causal model, and propose an extension of the ISR algorithm, called ISR-Multiclass.
This proposed algorithm can provably recover the invariant-feature subspace with ⌈ d spu /k
⌉ + 1 environments, where d spu is the number of spurious features and k is the number of classes.
Empirically, we first examine ISR-Multiclass in a synthetic dataset, and demonstrate its superiority over the original ISR in the multi-class setting.
Furthermore, we conduct experiments in Multiclass Coloured
MNIST, a semi-synthetic dataset with strong spurious correlations, and show that
ISR-Multiclass can significantly improve the robustness of neural nets trained by various methods (e.g., ERM and IRM) against spurious correlations.
Introduction
Domain generalization involves a learner having access to several domains during training time, which can be leveraged together to generalize better to unseen domains during test time.
Notably, a recent line of work (Wang et al.
[2022], Rosenfeld et al.
[2022], Kirichenko et al.
[2022]) focuses on achieving robustness to spurious correlations via simple linear transformations/last layer re-training.
In this work, we focus on and aim to extend the post-processing procedure proposed by Wang et al.
Wang et al.
[2022] present the ISR-Mean 1 , an algorithm that provably recovers the invariant feature subspace by using the 1 st order moments of the class conditional data distribution, under a causal model for classification considered by Rosenfeld et al.
[2022], Arjovsky et al.
However, in the more realistic multi-class setting, it is unclear whether there exists any relation between the number of classes and the number of training environments required to recover invariant features.
Further, Singla and Feizi [2021] demonstrates that features that are spurious for a given target label may be the core features for another.
One could apply class-specific transformations to mitigate this.
However, this is infeasible as we do not know the class labels during testing.
In light of the above challenges, we are specifically motivated by the following question:
In the context of linear causal models, can the number of classes compensate for the number of training environments in recovering the invariant features?
In order to answer the above questions, we propose ISR-Multiclass , an extension of the ISR-Mean algorithm that provably recovers the invariant subspace in ⌈ d spu /k
⌉ + 1 environments, where d spu is the number of spurious features (i.e., the dimensionality of the spurious-feature subspace) and k is the number of classes.
Note that our result improves over the original environmental complexity of d spu + 1 for binary classification problem.
2 , hence it provides an affirmative answer to the above problem.
Furthermore, our result shows that the benefits of k -class classification problems help to reduce the environmental complexity by an order of k .
Additionally, we introduce a new multi-class benchmark as a linear unit test based on Aubin et al.
[2021] and show that ISR-Multiclass can leverage class information to its advantage for faster recovery of invariant features.
We also demonstrate significant improvements in performance on the multi-class Coloured MNIST dataset (Ahuja et al.
A summary is shown in Figure 1.
ISR Multiclass
In this section, we introduce ISR-Multiclass i.e. I nvariant Feature S ubspace R ecovery for Multi
Class classification.
2.1 Setup
We study the causal model as shown in Figure 2, similar to that of Rosenfeld et al.
Let y be sampled from a prior distribution of labels { y 1 , y 2 , ...., y k } .
Let the dimension of the invariant and spurious features be d c and d s respectively, such that d c + d s = d , which is the total input dimension.
Then, we sample:
[Equation]
where µ y ∈ R d c , µ ye ∈ R d s and σ c , σ e ∈ R + .
As a comparison, in the original causal model for binary classification considered in Wang et al.
[2022], Rosenfeld et al.
[2020], the invariant feature distribution is given by z c ∼ N ( yµ c , σ 2 c I d c ) , where y ∈ { +1 , − 1 } .
We point out that this is an unnecessary assumption on the causal model mainly imposed for technical convenience.
Furthermore, as it will become clear shortly, this symmetry assumption on the means of the invariant feature distributions also artificially increases the environmental complexity from d spu / 2 + 1 to d spu + 1 .
Finally, the input x is generated from a linear transformation of the concatenated invariant and spurious features as follows:
[Equation]
As a consequence, the marginal distribution of x is the following:
[Equation]
Here, A = R d × d c , B = R d × d s are transformation matrices and R =
[ A, B ] ∈ R d × d .
In this work, we mainly focus on the design of the recovery algorithm itself, so for clarity of the derivation and discussion, we conveniently assume that the algorithm has access to infinite amount of data sampled according to the data model described above to limit the impact of finite samples, as a common practice also used in the prior works [Wang et al., 2022, Rosenfeld et al., 2020].
Similar to Wang et al.
[2022], we adopt the full-rank assumption on the means of the spurious features from different environments and different classes.
Assumption 1 For the set of environmental means, { µ ye } Ee =1 , ky =1 , we assume that
[Equation]
Intuitively, this assumption ensures that the spurious (and thus invariant) subspace can be recovered and multiple classes / multiple environments do not trivially replicate information.
Further, we also assume the following as per Rosenfeld et al.
Assumption 2 R is injective.
2.2 ISR-Multiclass
We now present ISR-Multiclass, as outlined in Algorithm 1.The detailed steps can be found in
Appendix A.1.
Briefly, consider the means ¯ x ek =
Aµ k + Bµ ke for a given class k and environment e .
Then, let:
[Equation]
Eigendecomposition of each M k results in E − 1 eigenvectors (from Assumption 1 and Wang et al.
Now, consider M total defined as:
[Equation]
where P i is the set of E − 1 eigenvectors obtained from the eigendecomposition of M k .
We now present Theorem 1 which formally states the proposed improvement.
The proof can be found in
Appendix A.2.
Theorem 1 (ISR-Multiclass) Assume that E ≥ ⌈
d s /k
+
and we have infinite data samples from every environment.
We perform SVD for M total (defined in (7) ), let { λ 1 , . . .
, λ d } denote the set of singular values obtained in descending order.
It is guaranteed that the top d s singular values are non-zero, i.e.,
[Equation]
The eigenvectors corresponding to these top d s singular values, i.e., { P ′ 1 , . .
, P ′ d s } , correspond to the d s spurious dimensions.
Consequently, the null space NullSpace ([ P ′ 1 , . .
, P ′
d s ] T ∈ R d s × d ) =
[ P ′′ 1 , . .
, P ′′ d c ] T ∈ R d c × d corresponds to the invariant dimension subspace.
A classifier f fitted to training data transformed by x (cid:55)→ [ P ′′ 1 , . . .
, P
[Table] [Table]
Optimality Our method involves applying an additional SVD operation over ISR-Mean, and thus the global optimality holds: a classifier trained on these features is globally optimal, similar to arguments in Wang et al.
Environment
Complexity The environment complexity of ISR-Multiclass is ⌈ d s /k
⌉ + 1 (detailed proof in A.2).
The key observation here is that the column space of each P i matrix for i ∈
[ k ] only consists of the span of { µ ie } Ee =1 .
So in order to identify the subspace of spurious features, one needs at least d s linearly independent components.
Hence, we only need to ensure that ( E − 1) k ≥
d s so that the column space of M total is full rank.
Solving this inequality leads us to the desired bound on the environmental complexity.
This implies that we can leverage information from multiple classes to reduce the environmental complexity by a factor of 1 /k , as compared to d s + 1 proposed in
Wang et al.
[2022], while only relying on the 1 st order moments of the data generating distribution.
Intuitively, our method leverages additional information from multiple classes to find the common spurious feature subspace in lesser number of environments.
Experiments
We now present empirical improvements demonstrated by ISR-Multiclass on 2 datasets.
[Figure]
3.1 Multi-Class Linear Unit Test
We construct a multi-class version of Example 3 (and its scrabled version, Example 3s) as used in
Aubin et al.
This is a synthetic dataset based on our causal model in section 2.1.
Specific details of our construction can be found in Appendix B. Figure 3 depicts the improvement of ISR-
Multiclass compared to the original ISR-Mean 3 and ERM Vapnik [1991].
Evaluation is performed on the test split where the spurious dimensions are randomized.
The Oracle is trained on this test split.
From figure 3, we observe that ISR-Multiclass is indeed able to leverage class information and recover invariant features to achieve optimal error with the number of environments inversely proportional to k , confirmed by our theoretical claim.
Especially in the last three plots - with greater classes and lesser environments ( k > 5 and n env < 5 ), both ISR-Mean and ERM incur higher error, but
ISR-Multiclass takes advantage of multiple classes to improve its performance instead and match the oracle.
3.2 Multi-Class Coloured MNIST (MC-CMNIST)
We consider the 10-class classification task of Coloured MNIST as proposed in Ahuja et al.
Note that in the train environments, every digit is highly correlated with a specific color, as depicted in
Figure 1.
This correlation breaks down in the test environment, i.e., every digit is randomly colored.
Evaluation Table 1 presents the results on MC-CMNIST.
Performance is evaluated on every group, which denotes a specific combination of ( y, color ) ∈
G =
Y × E .
Note that there are 10 × 10 =groups on this dataset.
During training, the samples across input groups are imbalanced owing to the spurious correlation where every digit majorly occurs in its associated color.
During test, samples across groups become balanced - thus testing a method’s ability to generalize to minority groups existing in the training set.
[Table]
We report the average accuracy (across all groups), worst group accuracy, and worst-10 group accuracy (average across 10 worst groups).
We compare the performance of ISR-Multiclass with
ERM, IB-ERM, IRM 4 and IB-IRM which are proposed by Arjovsky et al.
[2019] and Ahuja et al.
The oracle on this dataset achieves 99.03 ± 0.08 average accuracy as per Ahuja et al.
More details can be found in Appendix B.
It is evident that post-processing with ISR-Multiclass significantly improves both the average and worst-group accuracies, especially prevalent for ERM and IB-ERM.
While IRM and IB-IRM perform better than their ERM counterparts, ISR-Multiclass still improves the accuracy by ≈ 5 − 10% .
Note that ISR-Multiclass is a simple post-processing technique as compared to other methods which rely on end-to-end training.
We discuss the merits of this in more detail in Appendix C.
Conclusion
In this work, we propose ISR-Multiclass, an extension of ISR-Mean to a more practical setting of multi-class classification.
We theoretically prove that ISR-Multiclass can recover invariant features in fewer environments by using class information, specifically in ⌈ d spu /k
⌉ + 1 environments.
We corroborate our theory with empirical results by curating a new multi-class linear unit test and demonstrate faster recovery of invariant features by ISR-Multiclass.
Further, our method outperforms current methods on the MC-CMNIST by a large margin, both for average and worst group accuracy.
A Appendix A.1 Detailed Steps for ISR-Mutliclass
Step 1.
Estimating sample means for every environment and every class
Construct the following matrix M k for class k where each row contains the sample mean conditioned on a given environment e , for class k .
In other words, each row is ¯ x ek =
Aµ k + Bµ ke .
Note that in the infinite sample setting considered in our work, this is exactly the mean as per Equation (4).
[Equation]
Step
. Eigendecompose every
M k
For every class, eigen-decompose to obtain eigenvectors { P i }
di ∈ d corresponding to eigenvalues { λ i } di =1 .
By assumption 1, we obtain E − 1 eigenvectors (-1 from the mean centering in PCA) which correspond to non-zero eigenvalues.
[Equation]
where P ik is the i th eigenvector corresponding to a non-zero eigenvalue in the phase transition of M k .
Thus, P k recovers spurious dimensions corresponding to class k , as follows from ISR-Mean in Wang et al.
Note that unlike Wang et al.
[2022], we do not impose the condition that E > d s as information from a single class may not be sufficient to recover all spurious (and thus invariant) features.
Step 3:
Stack all P k and take SVD (Singular Value Decomposition)
After obtaining P k for every class y k , we stack all P k to obtain a new matrix M total as follows:
[Equation]
Next, take the SVD of P k .
Note that this step is motivated by the flag-mean Marrinan et al.
[2014] to find the common spurious subspace for all class labels.
Step 4: Extract Spurious Feature Subspace In Theorem 1, we prove that SVD of M total leads to d s non zero singular values, where the corresponding vectors { P ′ 1 , . . .
, P ′ d s } recover the underlying spurious subspace.
We stack these vectors as a matrix P ′ .
[Equation]
Step 5: Train a Classifier in the Null Space of Spurious-Feature Subspace This final step involves training a classifier in the null space of the extracted spurious feature subspace, which is the invariant feature subspace:
[Equation]
A.2 Proof for Theorem
Consider the matrix M total as defined in (10):
[Equation]
[Equation]
By definition, rank ( M total ) ≤ min ( d, ( E − 1) × k ) .
This trivially implies that: rank ( M total ) ≤ ( E − 1) × k .
Recall that each P k recovers the spurious dimension specific to class k .
In order to recover the underlying d s dimensional subspace, the rank of M total =
d s .
Combining this fact with the above statement, we obtain the following inequality:
[Equation]
Thus, the minimum number of environments required to recover the spurious (thus invariant) feature subspace benefits by leveraging information from classes.
The greater the number of classes, the lesser the number of training environments we require to recover the d c dimensional invariant features.
It should be noted that this decrease is observed while only leveraging the 1 st order moments of the class conditional data distribution.
Assuming condition (12) is satisfied, the rank of M total is capped at d s .
Thus, the SVD will lead to d s positive singular values.
We can then obtain the d s eigenvectors corresponding to these non-zero singular values, which span the spurious dimensions.
The transformation matrix will be P ′ ∈
R
d s × d .
Since z s ⊥ z c as per the setup 4, the null space of P ′ will correspond to vectors spanning the d − d s = d c dimensions as follows:
[Equation]
Finally, training on this invariant subspace will help us obtain the optimal invariant predictor as defined in Rosenfeld et al.
[2020], which completes the proof.
B Experimental Details B.1 Multiclass Linear Unit Test
We now discuss experimental details for constructing and evaluating on Example3/3s-Multiclass, a
Multi Class Linear Unit Test.
Construction First, we sample y from a multinomial distribution of uniform probability 1 /k , where k is the number of classes.
Then, the first d c invariant features are sampled from a Gaussian distribution where the mean depends on the class label.
Similarly, the next d s spurious features are sampled from a Gaussian distribution where the mean now depends on the class label as well as the environment label.
This can be formulated as follows:
For a given environment e ,
[Equation]
For Example3-Multiclass, R = I d , and for Example3s-Multiclass (scrambled variation), R ∈ R d × d is an orthonormal matrix.
The sampling of means µ k , µ ke is done from a uniform distribution between [0 , 1) .
ν inv and ν spu are the scale of invariant and spurious features.
We set µ inv = 0 .
1 and µ spu = 1 as regularization may encourage learning spurious features, making it harder to learn the invariant features.
This is similar to Example2 in Aubin et al.
In our experiments, σ c = 0 .
1 , σ e = 0 .
1 , d s = 5 , d c = 5 .
We sample 10,000 points per environment.
k varies from 2 to 7 .
Code and Hyperparameters For all methods, we perform a hyperparameter search over 5 data seeds and 5 model trials.
In every trial, we train the algorithm on the train split and use the Adam Kingma and Ba [2014] optimizer for optimization.
The model with the least mean validation error across all environments is chosen.
For ISR-Mean, we use the implementation from their released code at https://github.com/Haoxiang-Wang/ISR .
B.2 Multiclass Coloured MNIST (MC-CMNIST)
Figure B.2 depicts a summary of the dataset - for every digit, the corresponding colour is highly correlated in the training set.
This correlation breaks during testing.
We directly employ the Multiclass coloured MNIST dataset, models and hyperparameters provided by Ahuja et al.
[2021] at https://github.com/ahujak/IB-IRM .
For ERM, IB-ERM, IRM and IB-IRM, we run a sweep over hyperparameters using the grid as suggested above.
The best model is chosen by using train domain validation (Gulrajani and LopezPaz [2020]).
ISR-Multiclass is applied on the last-layer over the classification weights to enable the invariant feature subspace transformation.
Note that ISR-Multiclass uses colour labels as the group information, and we ensure this same definition applies to IRM to ensure a fair comparison.
It should also be noted that similar to Wang et al.
[2022], we adopt a strategy of scaling down the spurious dimensions extracted from convolutional networks.
C Post-processing v/s end-to-end training
It should be noted that ISR-Multiclass is a post-processing technique and can be applied on top of any pretrained embeddings, while other methods rely on end-to-end training.
This further demonstrates the possibility of using such a linear transformation on any embeddings obtained from large pretrained models.
Fine-tuning such models maybe infeasible, and applying ISR-Multiclass (vis-a-vis linear probing) can give us the additional benefit of robustness to spurious correlations.