Feature-Refinement Classification Head (FCH) for Modeling Inter-Sample Relationships in Neural Networks

Abstract

Conventional classification methods rely on the assumption that samples are independent and identically distributed, often ignoring the latent relational structures governing real-world data. Consequently, these methods fail to capture cross-sample dependencies, resulting in impoverished feature representations and suboptimal generalization. To address this limitation, we propose the Feature-refinement Classification Head (FCH), a novel and modular component that explicitly models inter-sample relationships within training batches. FCH constructs an adjacency matrix from embedding vectors and employs a graph neural network to refine features via relational propagation. Seamlessly integrated and architecture-agnostic, FCH leverages the end-to-end optimization using cross-entropy and structure preserving loss as joint objectives. Extensive experiments on benchmarks such as CIFAR-10, MNIST, and STL-10 demonstrate that FCH consistently enhances accuracy, precision, recall, and F1-scores across diverse backbones. Notably, FCH demonstrates significant performance improvements even in lightweight models, demonstrating its robustness, scalability, and practicality. Overall, FCH offers a principled, generalizable alternative to traditional classification heads by effectively leveraging sample-to-sample relationships.