[Paper Review] BatchFormer: Learning to Explore Sample Relationships for Robust Representation Learning

Zhi Hou, Baosheng Yu+1 more|arXiv (Cornell University)|Mar 3, 2022

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

BatchFormer introduces a plug-and-play transformer module that operates on the batch dimension to learn sample relationships during training, enabling cross-sample knowledge transfer—especially benefiting low-shot and tail-class recognition. By sharing a classifier before and after the module, it eliminates inference overhead and achieves state-of-the-art performance across long-tailed recognition, zero-shot learning, domain generalization, and contrastive learning without additional components.

ABSTRACT

Despite the success of deep neural networks, there are still many challenges in deep representation learning due to the data scarcity issues such as data imbalance, unseen distribution, and domain shift. To address the above-mentioned issues, a variety of methods have been devised to explore the sample relationships in a vanilla way (i.e., from the perspectives of either the input or the loss function), failing to explore the internal structure of deep neural networks for learning with sample relationships. Inspired by this, we propose to enable deep neural networks themselves with the ability to learn the sample relationships from each mini-batch. Specifically, we introduce a batch transformer module or BatchFormer, which is then applied into the batch dimension of each mini-batch to implicitly explore sample relationships during training. By doing this, the proposed method enables the collaboration of different samples, e.g., the head-class samples can also contribute to the learning of the tail classes for long-tailed recognition. Furthermore, to mitigate the gap between training and testing, we share the classifier between with or without the BatchFormer during training, which can thus be removed during testing. We perform extensive experiments on over ten datasets and the proposed method achieves significant improvements on different data scarcity applications without any bells and whistles, including the tasks of long-tailed recognition, compositional zero-shot learning, domain generalization, and contrastive learning. Code will be made publicly available at https://github.com/zhihou7/BatchFormer.

Motivation & Objective

• To address data scarcity challenges in deep representation learning, such as long-tailed class distributions, domain shift, and unseen classes.
• To overcome the limitation of existing methods that explore sample relationships only via input or loss functions, rather than through the internal structure of deep neural networks.
• To enable deep neural networks to implicitly learn and exploit relationships among samples in each mini-batch during training.
• To reduce the training-inference gap by sharing the classifier before and after the BatchFormer module, allowing removal of the module at test time.

Proposed method

• Introduces a BatchFormer module that treats each sample in a mini-batch as a token in a sequence, applying a standard Transformer encoder to model inter-sample relationships.
• Applies the BatchFormer to the batch dimension, enabling gradient flow from each sample's loss to all other samples in the mini-batch, effectively creating virtual data augmentation.
• Uses a shared classifier before and after the BatchFormer during training to enforce batch-invariant representation learning and allow removal of the module at inference time.
• Employs balanced softmax and contrastive learning objectives to enhance feature discrimination and generalization, particularly for rare classes.
• Leverages Grad-CAM and gradient analysis to visualize how BatchFormer improves attention on relevant object parts and reduces spurious correlations.
• Conducts extensive ablation studies on batch size, classifier sharing, and backbone architectures to validate design choices and robustness.

Experimental results

Research questions

• RQ1 Can modeling inter-sample relationships within a mini-batch improve representation learning under data scarcity?
• RQ2 Does a batch-wise transformer module enhance feature learning for low-shot and tail classes in long-tailed recognition?
• RQ3 Can the BatchFormer module be integrated into existing models without modifying inference pipelines?
• RQ4 How does the shared classifier mechanism mitigate the training-inference gap in batch-based representation learning?
• RQ5 To what extent does BatchFormer improve robustness in zero-shot learning, domain generalization, and contrastive learning?

Key findings

• BatchFormer achieves state-of-the-art performance on over ten datasets across multiple data-scarce tasks, including long-tailed recognition, zero-shot learning, domain generalization, and contrastive learning.
• On ImageNet-LT, BatchFormer achieves 50.9% top-1 accuracy on the 'All' class, outperforming the baseline with shared classifier and significantly improving few-shot performance.
• The model with BatchFormer converges faster, reaching moco-v3-level performance in only 38 epochs, demonstrating improved training efficiency.
• Abllation studies show that BatchFormer is less sensitive to batch size when below 128, and batch size 512 yields the best performance on few-shot categories.
• Grad-CAM visualization confirms that BatchFormer enhances attention on object parts and suppresses spurious background correlations, especially in complex scenes.
• Gradient analysis reveals that rarer classes exert larger gradients on other images in the mini-batch, indicating implicit data augmentation that benefits low-shot learning.