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 Invariantfeature Subspace Recovery (ISR), a domain generalization algorithm which uses the means of class-conditional data distributions to provably identify the invariantfeature subspace.
However, the original ISR algorithm is conditioned on single class only, without utilizing information from the rest classes.
In this work, we consider the setting of multi-class classification, 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 nmnumber 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 the post-processing procedure proposed by Wang et al.
They present 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 binary data model.
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 these, we are specifically motivated by the following question:
Can information from multiple classes help us recover the invariant-feature subspace with fewer environments while only using the 1 st order moments of the data distribution?
In order to answer the above questions, we propose ISR-Multiclass , an extension of the ISR-Mean algorithm which provably recovers the invariant subspace in O ( d spu /k ) 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.
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.1, similar to 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 + .
Finally, the input x is sampled from the following distribution:
[Equation]
Here, A = R d × d c , B = R d × d s are transformation matrices and R =
[ A, B ] ∈ R d × d .
Note that we assume access to infinite-samples in our case.
We adopt a degeneracy assumption with multiple classes stated as follows:
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: the algorithm and 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.3.
Theorem 1 (ISR-Multiclass) Let us assume that E ≥ d s /k
+
and we have infinite data samples from every environment.
We perform SVD for M total (defined in (5) ), 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]
≤ d s , λ i >
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 d c ] T x recovers the invariant optimal predictor.
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.
[2022] and Rosenfeld et al.
Environment Complexity
The environment complexity of ISR-Multiclass is O
( d s /k )
(A.3).
This implies that we can leverage information from multiple classes to reduce the environmental complexity by a factor of 1 /k , as compared to O ( d s ) 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] [Figure]
M e a n E rr o r
Num Classes:   8 10 n_env
Num Classes:   8 10 n_env Num Classes:   8 10 n_env
Num Classes:   8 10 n_env Num Classes:   8 10 n_env
Num Classes:
ERM
ISR-mean
ISR-MultiClass
Oracle
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 ISRMulticlass compared to the original ISR-Mean 2 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
[Table]
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]
average and worst group accuracies, especially for ERM.
Note that the variance for worst group accuracies is high because of the less number of samples per group ( ≈ 300 ).
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 3 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 ISR-Multiclass Algorithm
Algorithm 1 presents the steps used for applying ISR-Multiclass.
A.2 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 .
[Equation]
x T 1 k .
[Equation]
 (6)
Step 2.
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]
(7) 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.3 Proof for Theorem
Consider the matrix M total as defined in (8):
[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 2, 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 , y ∼ Multinomial (cid:18)k (cid:19) , z c ∼ { N ( µ k , σ c I d c ) ∗ ν inv for y = k, z e ∼ { N ( µ ke , σ e I d s ) ∗ ν spu for y
= k, env = e, z ← (cid:20)
z
c z e (cid:21) , x = Rz
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.