Take One Gram of Neural Features,
Get Enhanced Group Robustness

Previous work [21] observed that ERM tends to fit models on data presenting easy-to-learn spurious correlations at the beginning of the learning process. We hence train for a few iterations an exogeneous “identification model” –a convolutional neural network $\Phi$ composed of $L$ layers with parameters $\omega$, pre-trained on ImageNet [7]– by empirical loss minimization:

where the cross-entropy loss $\ell$ is applied between the model’s prediction $\Phi (x_i,\omega )$ and the true label $y_i$ for sample $x_i$. After this initial training, and in the rest of the paper, the parameters $\omega$ of the identification model $\Phi$ are frozen.

We now turn to feature-based clustering. We denote the feature map of an image $x$ at layer $l$ of $\Phi$ by ${\varphi }_l\left(x\right)\in R^{M_l\times C_l}$, where $C_l$ is the number of channels and $M_l$ is the spatial dimension of the feature map. For each image $x_i$, we extract its feature maps at $S⩽L$ different and fixed layers $l_1,\dots ,l_S$, and compute the Gram matrices defined as:

We then vectorize and normalize each of these $S$ Gram matrices:

The normalization permits us to balance evenly the contribution of each Gram matrix in the clustering loss. The “environment” of each image $x_i$ is thus encoded by the vector $f_i=[f_{i,l_1};\cdots ;f_{i,l_S}]\in R^D$, where $D={\sum }_{l=l_1}^{l_S}{C}_l^2$. We discover $E^{\prime }$ environments by clustering the $N$ training images into $E^{\prime }$ clusters $C_1,\dots ,C_{E^{\prime }}$, via $k$-means clustering, i.e., by computing a solution to:

where ${\Vert f_i-f_j\Vert }_2^2={\sum }_{l=l_1}^{l_S}{\Vert f_{i,l}-f_{j,l}\Vert }_2^2$.

To overcome the computational cost of storing all these vectors and computing distances between them in high dimension, we perform random projections of the vectors $f_{i,l}$ in a lower-dimensional space as proposed in [1] (see more details in the supplementary material)

As a results of the clustering, each training image is now equipped with a pseudo-environment label $\hat{e}\in \{1\cdots E^{\prime }\}$. Combined with its class label, this provides a pseudo-group label $\hat{g}=(\hat{e},y)$. The training set being now partitioned into pseudo-groups, we train a robust classifier $h$, distinct from $\Phi$, with parameters $\theta$, by minimizing the worst-group risk (GroupDRO [24]):

based on the cross-entropy loss $\ell$, where $D_{e,k}\subset D$ denotes the set of samples with pseudo-group label $\hat{g}=(e,k).$

Unlike previous approaches [6, 27, 2] that need true group labels in the validation set to define and assess worst-group performance as the metric to set hyperparameters, we first partition the validation set using the clusters found on the training set and then conduct cross-validation based on the resulting pseudo-groups. In our experiments, we observe that this type of model selection is effective to achieve proper group robustness.

We experiment with three standard image classification datasets on which previous works evaluate worst-group performance: Waterbirds [24] is a dataset composed of bird photographs from the CUB dataset [29] superimposed on background scenes taken from the Places365 dataset [31]. The target labels are “landbird” and “waterbird” which are spuriously correlated with the background images of either “land” or “water”. We evaluate on the test set with the average accuracy and the worst-group accuracy (“waterbird” on “land”); CelebA [22] is a celebrity large-scale face dataset with 202,599 natural images. There exists a spurious correlation between the hair color and the gender (“male” or “female”) of a person. We evaluate on the test set with the average accuracy and the worst-group accuracy (“male” with “blond” hair); COCO-on-Places-224 is the same dataset as in [2] but with images resized to $224\times 224$ (instead of $64\times 64$ in the original paper). There are 10 segmented COCO [20] objects superimposed on scenes from the Places365 dataset. This time, a group of backgrounds are spuriously correlated with each object at training time. We evaluate the accuracy on a first test set with objects on the same backgrounds as during training, called the in-distribution set (‘ind’), and on a second test set with objects on unseen backgrounds, dubbed the systematically-shifted set (‘sys’).

We compare our approach against the standard ERM baseline and recent methods that aim at robust predictions across groups without the use of train group annotations (EIIL [6], GEORGE [27] and JTT [21]). We also include robust methods that use true group annotations at train time (IRM [3], importance weighting and GroupDRO [24]). The latter methods and ERM were already implemented and we took care to reproduce results for all methods. Note that our approach and GroupDRO share the same robust optimization objective.

All methods use a ResNet-50 architecture pre-trained on ImageNet [7] as the robust classifier (classifier $h$ in Section 2). Models are optimized using SGD with momentum. For GroupDRO and ERM, we use the hyperparameters reported by the authors on Waterbirds and CelebA datasets. Note that hyperparameters have been selected with the use of a validation set with group labels. Regarding our approach, we select a VGG-19 [26] architecture for the identification model ($\Phi$ in Section 2) and train it for $1$ epoch using SGD with momentum. Among usual layers used to compute style representations in neural style transfer, we observed improved performance by selecting deeper layers in the network: for each dataset, we consistently extract features from the conv5   1 layer, i.e., the first convolutional layer of block $5$. We include results with two types of model: (i) based on validation set with true-group annotations (‘GramClust-orig’); (ii) based on pseudo-group labels (‘GramClust-cv’) predicted by our clustering.

We report quantitative comparisons on Waterbirds, CelebA and COCO-on-Places-224 in Table 1. First, we observe that GramClust improves worst-group test accuracy over ERM baseline on Waterbirds and CelebA and sys accuracy on COCO-on-Places224. More importantly, GramClust-cv achieves state-of-the-art performances on group robustness compared to all methods that do not use group labels on the training set. On CelebA, which is a large-scale datasets with natural images, our approach outperforms the previous best, JTT [21], by $21.9$ pts. We were not able to scale EIIL [6] on this dataset due to memory overflow issues. Note that GramClust-orig uses the same hyperparameters as EIIL, GEORGE and JTT for robust training of the target classifier from predicted group labels, and still displays significant improvements. Surprisingly, GramClust-cv and GramClust-orig outperform GroupDRO on Waterbirds with $85.3\%$ vs. $83.9\%$, while the latter method uses true-group labels during training. This may be due to the ambiguity of the background in some Waterbirds images.

Since [15] uses the channel-wise mean and variance of image features to perform style transfer, we compare the use of such style statistics (‘MeanVar’) against our use of Gram matrices in Table 2. MeanVar reaches test worst-group accuracy on-par with Gram matrix on Waterbirds but degrades significantly performances on CelebA. Gram matrices provide more information than MeanVar as their diagonals already contain the information about the channel-wise mean and variance of the deep features (see Eq. 2). Hence, these results show that, when scaling on large and natural-image datasets such as CelebA, keeping all the correlations between different channels is important for group robustness.

We also compared our use of VGG-19 conv5   1 features to capture style with the direct use of the penultimate (‘AvgPool’) representation of a more modern ResNet-50 identification model. Note that, while dating back to 2015, VGG features are still successfully used through their Gram matrices, e.g. in [14, 18, 5]. In Table 2, we observe that using the penultimate layer of a ResNet-50 as style representation for the clustering produces poorer performance.

We study the behavior of our clustering algorithm w.r.t. the layers selected to extract features and to the number of clusters. This analysis is conducted on the Waterbirds dataset.

First, we evaluate the impact of the selection of VGG-19 layers to extract the features in the clustering stage. To this end, we study the matching of the predicted environments to the true environment labels on the validation via Hungarian matching [17] and measure the global matching accuracy across all validation samples for each five layers commonly used in neural style transfer (conv1   1,conv2   1,conv3   1,conv4   1,conv5   1). Results in Figure 3 show that: Features from deeper layers correlate with better matching accuracy; Our approach is robust to the choice of deep layers either taken together (allconvX   1) or individually such as conv4   1 and conv5   1; Using conv5   1 outperforms selecting all traditional style layers. Consistent conclusions are found on the CelebA dataset (see Supplementary material).

Second, we study the impact of the number of clusters as hyperparameter in the clustering algorithm. Worst-group accuracy on the validation set for $E^{\prime }\in \{2,4,8,16,32\}$ clusters are reported in Figure 3. Or method is robust to a variation in the number of clusters: GramClust with more clusters than actual environments produces a slight drop in performance but still improves performance over ERM and remains on-par performances with GroupDRO.