Get the Best of Both Worlds: Improving Accuracy and Transferability by Grassmann Class Representation

TL;DR

This paper introduces Grassmann Class Representation (GCR), which generalizes class vectors to subspaces on the Grassmann manifold, improving neural network accuracy and transferability by jointly optimizing class subspaces with deep models.

Contribution

The paper proposes a novel GCR method that models classes as subspaces, enhancing representation power and transferability, and integrates Riemannian SGD into deep learning frameworks.

Findings

Significant error reduction on ImageNet-1K across multiple architectures.

Improved transfer accuracy on downstream tasks with GCR.

Increased intra-class feature variability with larger subspaces.

Abstract

We generalize the class vectors found in neural networks to linear subspaces (i.e.~points in the Grassmann manifold) and show that the Grassmann Class Representation (GCR) enables the simultaneous improvement in accuracy and feature transferability. In GCR, each class is a subspace and the logit is defined as the norm of the projection of a feature onto the class subspace. We integrate Riemannian SGD into deep learning frameworks such that class subspaces in a Grassmannian are jointly optimized with the rest model parameters. Compared to the vector form, the representative capability of subspaces is more powerful. We show that on ImageNet-1K, the top-1 error of ResNet50-D, ResNeXt50, Swin-T and Deit3-S are reduced by 5.6%, 4.5%, 3.0% and 3.5%, respectively. Subspaces also provide freedom for features to vary and we observed that the intra-class feature variability grows when the subspace dimension increases. Consequently, we found the quality of GCR features is better for downstream tasks. For ResNet50-D, the average linear transfer accuracy across 6 datasets improves from 77.98% to 79.70% compared to the strong baseline of vanilla softmax. For Swin-T, it improves from 81.5% to 83.4% and for Deit3, it improves from 73.8% to 81.4%. With these encouraging results, we believe that more applications could benefit from the Grassmann class representation. Code is released at https://github.com/innerlee/GCR.