LossTransform: Reformulating the Loss Function for Contrastive Learning

Abstract

Contrastive learning improves model performance by differentiating between positive and negative sample pairs. However, its application is primarily confined to classification tasks, facing challenges with complex recognition tasks such as object detection and segmentation due to its limited capacity to capture spatial relationships and fine-grained features. To address this limitation, we propose LossTransform, a novel approach that redefines positive sample pairs and establishes a novel contrastive loss paradigm. LossTransform advances contrastive learning to the instance level, departing from the traditional sample level. Empirical evaluations on ImageNet, CIFAR, and object detection benchmarks indicate that LossTransform improves accuracy by +2.73% on CIFAR, +2.52% on ImageNet, and up to +5.2% in average precision on detection tasks, while maintaining efficiency. These results illustrate that LossTransform is compatible with large-scale training pipelines and exhibits robust performance across diverse and complex datasets. By optimizing model performance and significantly reducing training time, this research enables more efficient and accessible solutions for societal applications.

1. Introduction

Contrastive learning is instrumental in interpretable artificial intelligence as it enhances model performance by distinguishing between positive and negative pairs of samples. However, its practical applicability is limited. It performs well primarily in basic classification tasks on small datasets but struggles with more complex recognition tasks on larger datasets. Additionally, current studies [1,2] show that contrastive learning methods are not user-friendly, as they often require modifications to the original data pipeline, model architecture, and training recipe, instead of being directly applicable to existing models.

The application of contrastive learning is mainly limited to classification tasks due to its fundamental mechanism, which relies on unsupervised clustering, which brings positive sample pairs closer in the projection space while simultaneously maximizing the feature distance between negative sample pairs. Conventionally, the determination of whether two samples from a positive pair is based on two criteria: whether they originate from distinct augmentations of the same sample or whether they are separate sharing the same class label. However, this paradigm encounters challenges in more complex visual tasks such as object detection and semantic segmentation. These tasks not only involve category classification but also require the model to accurately localize objects spatially and capture fine-grained pixel-level information. The limitations of contrastive learning in these tasks arise from its inability to effectively manage multiple object instances and leverage local feature information.

To effectively apply contrastive learning in general recognition tasks, we have redefined the concept of positive sample pairs and designed a novel contrastive loss paradigm. Specifically, in object detection applications, our approach extends contrastive learning to facilitate object pairing at the instance level (in this context, a “sample” refers to a single input data point in general recognition tasks, while an “instance” is a term specific to object detection, where a sample image may contain multiple instances, each defined by a bounding box, label, and polygon boundary or masked area), thereby diverging from the conventional use of positive/negative image pairs. As shown in Figure 1, an image undergoes independent random data augmentations (e.g., flipping, translation, and cropping), resulting in multiple replicas. These replicas contain both common instances and potentially unique instances. The feature representation of an instance’s bounding box is subject to change after data augmentation. To align these common instances, we first apply an inverse data augmentation to restore these features to a unified space. Based on a self-supervised alignment mechanism, we use a greedy matching algorithm to subsequently select high-confidence positive instance pairs from the multiple image replicas. Finally, we minimize the Kullback–Leibler divergence to constrain the distance between these positive instance pairs, thereby ensuring that the representations of identical instances from different image replicas become more compact within the projection space.

Figure 1. Contrastive learning applied to object detection. Image replicas $R=\{{\psi }_1\left(x\right),{\psi }_2\left(x\right),\cdots \}$ are generated through Random Augmentation Sampling (RASping, see Section 2.3), where each replica ${\psi }_k\left(x\right)$ produces a series of bounding box $p_k=\{{\mathrm{bbox}}_k^1,{\mathrm{bbox}}_k^2,\dots \}$ via predictions. Inverse transformations (see Section 3.1) are applied to these boxes to estimate the model’s predictions on the original image. Geometric and statistical methods are used to align instances and match positive pairs of boxes (a bounding box labeled “✓” indicates that the instance can be paired with a bounding box in another image; unpaired bounding boxes are labeled “✗”, see Section 3.2), followed by the computation of an anchor box by averaging the corresponding box pairs. The KL divergence is minimized to reduce the distance between each predicted box and its anchor (see Section 3.3). Finally, anchor boxes are transformed and concatenated with unpaired bounding boxes to obtain the optimized box predictions as $P=\{{\psi }_1\left(\overline{p_i}\right),{\psi }_2\left(\overline{p_i}\right),\cdots \}$.

In this study, we conduct extensive evaluations of LossTransform on various datasets, model architectures, and recognition tasks. Our ablation studies show that LossTransform enables models to achieve improvements during testing: a +2.73% accuracy increase on CIFAR [3] and a +2.52% accuracy increase on ImageNet [4]. Experiments using the multi-task loss functions of Faster R-CNN [5] and Mask R-CNN [6] produce similar results, with a +5.2% Average Precision (AP) improvement on PennFudanPed [7], a +0.8% mean Average Precision (mAP) increase on PASCAL VOC 2012 [8], and a +0.4% mAP increase on MS COCO [9].

2. Related Work

2.1. Test-Time Augmentation

The ensemble learning method [10] and test time augmentation (TTA) technique [11] are often used in evaluating neural networks for image classification [12,13,14]. These methods aggregate outputs from various models or data augmentations and compute the average predicted results as the final prediction. From the perspective of knowledge distillation [15], both methods serve as strategies for generating soft labels [16]. TTA is inherently more versatile than other ensemble voting methods in a production environment because it does not require multiple models during the prediction phase. Recently, researchers applied this method to post-processing the prediction and retraining models [15,17,18]. We have incorporated this method in the training process to minimize the gap between predictions and the anchor prediction.

2.2. Semi-Supervised Learning

Different from data augmentation techniques, Semi-Supervised Learning (SSL) extends the cardinality of a training dataset by using a large amount of extra unlabelled data and subsequently generating pseudo labels through clustering algorithms [1,19]. In the SSL framework, the model can estimate the data-generating distribution and construct decision boundaries more precisely. An important concept in SSL is providing pseudo-labels to unlabeled data [20]. Building on this idea, in our method, if any images are successfully matched, they are collectively treated as positive samples of an unknown anchor sample. There is no requirement to identify the label of this unknown anchor sample; rather, we compute the average of the labels of all positive samples to serve as the pseudo-label for this anchor sample.

2.3. Random Augmentation Sampler

RASampler [21] is designed to limit data loading to a subset of the dataset for distribution purposes, incorporating repeated augmentation. For example, consider a dataset with five samples. To increase training data diversity, each sample can be augmented multiple times—here, three times per sample. The process typically involves

• Shuffling: Shuffling the original dataset indices (e.g., [0, 1, 2, 3, 4]) to create a randomized order (e.g., [3, 1, 4, 0, 2]).
• Repetition: Repeating each index based on the augmentation count (e.g., three times) results in [3, 3, 3, 1, 1, 1, 4, 4, 4, 0, 0, 0, 2, 2, 2].
• Subsampling: Splitting the indices to create a subset (e.g., [3, 3, 1, 4, 4]).

We use RAsamper in LossTransform to ensure that examples are repeated multiple times in each iteration while the total number of examples seen by the model in each epoch remains unchanged.

2.4. Sample Pairing in Object Detection

Sample pairing is a key issue when applying contrastive learning to object detection. Different works have proposed distinct strategies, but each of them comes with specific limitations.

2.4.1. Contrastive Learning for Object Detection

In this work, an image is augmented twice independently with random transformations, where random cropping plays a central role. The authors name this process TwoCropTransform [22]. It is important to note that this does not mean cropping out individual objects, but rather applying random cropping twice to generate two augmented views as a positive pair. This strategy may cause a mismatch in annotations: for example, when one crop removes objects on the left edge and the other crop removes objects on the right edge, the two augmented views will not share the same set of objects, resulting in unpaired labels and bounding boxes. In addition, the official code applies random horizontal flipping but does not introduce a corresponding alignment step in the loss function, which further increases the risk of annotation mismatch. Different from this work, our method does not rely on image-to-image positive pairs, but instead defines a pairing process at the object-to-object level.

2.4.2. DetCo: Unsupervised Contrastive Learning for Object Detection

DetCo [23] extends the previous idea by introducing two extra components. Besides random cropping, it generates small patches from the original image and further designs a global-to-global, global-to-local, and local-to-local pairing strategy. The final loss function is the weighted sum of these three terms. However, DetCo may face a similar issue as the first work. The paper does not clarify whether the crops are applied to the whole image or to detected objects. If the crops are image-level, the same annotation mismatch problem can still happen. Moreover, the released code only supports testing with pretrained models, making it difficult to confirm how sample pairing is implemented during training.

2.4.3. Contrastive Learning-Based Robust Object Detection Under Smoky Conditions

This work addresses detection in smoky images and builds a new architecture trained with a contrastive loss [24]. Unlike the previous two methods, sample pairing here does not suffer from annotation mismatch. This is because, under the smoky condition setting, the authors avoid using random cropping, image translation, or horizontal flipping. All augmented images share exactly the same annotations, naturally forming valid positive pairs. However, in general object detection tasks, random cropping, shifting, and flipping are critical for improving robustness. Therefore, this method may not generalize well beyond the smoky condition scenario.

2.5. Object Detection Loss Functions

Loss functions are a fundamental component in object detection, as they directly influence both classification and localization accuracy. A recent survey [25] provides a comprehensive review of detection losses, underscoring the role of carefully designed objectives in enhancing discriminative power. Our work follows a similar perspective from the lens of contrastive learning: by redefining positive sample pairs and proposing a novel contrastive loss, we contribute to the broader discussion on loss design. In particular, our instance-level contrastive paradigm complements detection-oriented loss formulations, highlighting the importance of loss functions as a unifying mechanism across recognition tasks.

3. Methods

Previous studies have relied heavily on minimizing the distance between positive pairs while maximizing the distance between negative pairs to improve model performance. This study demonstrates that substantial improvements in model performance can be achieved by employing LossTransform without the need to maximize the distance between negative pairs. In the following sections, we will take the task of object detection as an example: (1) defining the inverse transform for predicting the bounding box prediction; (2) explaining the instance alignment algorithm, which identifies all positive instance pairs through a greedy algorithm; and (3) presenting LossTransform, which redefines the loss function in contrastive learning and extends its application to non-classification tasks, while maintaining efficacy in classification tasks.

3.1. Inverse Transform

While data augmentation methods do not change the category of an instance, they can significantly impact its bounding box specifically; when performing a horizontal flip (hflip) or translation on an image, the transformation of the bounding box $(c_x,c_y,h,w)$ follows these steps:

$$\begin{array}{cc}\hfill {\psi }_{\mathrm{hflip}}& :(c_x,c_y,h,w)\mapsto (W-c_x,c_y,h,w),\hfill \end{array}$$

(1)

$$\begin{array}{cc}\hfill {\psi }_{\mathrm{translation}}& :(c_x,c_y,h,w)\mapsto (c_x+t_x,c_y+t_y,h,w),\hfill \end{array}$$

(2)

$$\begin{array}{cc}\hfill {\psi }_{\mathrm{scaling}}& :(c_x,c_y,h,w)\mapsto (s_x{c}_x,s_y{c}_y,s_{y}h,s_{x}w),\hfill \end{array}$$

(3)

$${\psi }_{\mathrm{rotatation}}:\left[\begin{array}{c}{c}_x\\ c_y\\ h\\ w\end{array}\right]\mapsto \left[\begin{array}{cccc}cos\theta & -sin\theta & 0& 0\\ sin\theta & cos\theta & 0& 0\\ 0& 0& |cos\theta |& |sin\theta |\\ 0& 0& |sin\theta |& |cos\theta |\end{array}\right]\left[\begin{array}{c}{c}_x-W/2\\ c_y-H/2\\ h\\ w\end{array}\right]$$

(4)

where $c_x$ and $c_y$ are the center coordinates, $w,h$ are the width and height, respectively, $W,H$ are the image width and height, $t_x,t_y$ are the translation offsets along the horizontal (x-axis) and vertical (y-axis) directions, respectively, and $\theta$ is the rotation angle. Note that, for the object detection task, arbitrary-angle image rotation is not feasible because bounding boxes cannot be accurately computed unless the rotation angle $\theta \in \{0,\pi /2,\pi ,3\pi /2\}$. The inverse transformation is straightforward, as follows:

$$\begin{array}{cc}\hfill {\psi }_{\mathrm{hflip}}^{-1}& :(c_x^{\prime },c_y^{\prime },h^{\prime },w^{\prime })\mapsto (W-c_x^{\prime },c_y^{\prime },h^{\prime },w^{\prime }),\hfill \end{array}$$

(5)

$$\begin{array}{cc}\hfill {\psi }_{\mathrm{translation}}^{-1}& :(c_x^{\prime },c_y^{\prime },h^{\prime },w^{\prime })\mapsto (c_x^{\prime }-t_x,c_y^{\prime }-t_y,h^{\prime },w^{\prime }).\hfill \end{array}$$

(6)

$$\begin{array}{cc}\hfill {\psi }_{\mathrm{scaling}}^{-1}& :(c_x^{\prime },c_y^{\prime },h^{\prime },w^{\prime })\mapsto (c_x^{\prime }/s_x,c_y^{\prime }/s_y,h^{\prime }/s_y,w^{\prime }/s_x),\hfill \end{array}$$

(7)

$${\psi }_{\mathrm{rotatation}}^{-1}:\left[\begin{array}{c}{c}_x^{\prime }\\ c_y^{\prime }\\ h^{\prime }\\ w^{\prime }\end{array}\right]\mapsto \left[\begin{array}{cccc}cos\theta & sin\theta & 0& 0\\ -sin\theta & cos\theta & 0& 0\\ 0& 0& |cos\theta |& |sin\theta |\\ 0& 0& |sin\theta |& |cos\theta |\end{array}\right]\left[\begin{array}{c}{c}_x^{\prime }+W/2\\ c_y^{\prime }+H/2\\ h^{\prime }\\ w^{\prime }\end{array}\right]$$

(8)

3.2. Instance Alignment

In our method, we use random data augmentation to generate multiple replicas. Given an image with raw annotations $y=\{{\mathrm{bbox}}_1,{\mathrm{bbox}}_2,\dots \}$, generate m augmented replicas by applying a random data augmentation ${\psi }_j\in \{{\psi }_{\mathrm{hflip}},{\psi }_{\mathrm{translation}},\cdots \}$ to the raw image, producing corresponding predictions

$${\psi }_j\left(p\right)=\{{\mathrm{bbox}}_1^j,{\mathrm{bbox}}_2^j,\dots \},j=1,\cdots m.$$

(9)

We align the predicted bounding boxes for positive instances by sorting them according to

$${\mathrm{bbox}}_i^j=\underset{\mathrm{bbox}\in {\psi }_j\left(p\right)}{argmax}\mathrm{IoU}\left({\mathrm{bbox}}_i,{\psi }_j^{-1}\left(\mathrm{bbox}\right)\right).$$

(10)

Now, the averaged prediction is defined as $\overline{p}=\{\overline{{\mathrm{bbox}}_1},\overline{{\mathrm{bbox}}_2},\dots \}$. The aligned bounding boxes of instances are given by (see Figure A2 for visualization)

$$\overline{{\mathrm{bbox}}_i}={\displaystyle \frac{1}{m}}\sum _j{\psi }_j^{-1}\left({\mathrm{bbox}}_i^j\right).$$

(11)

Additionally, for classification tasks, the ground-truth label is not affected by data augmentation. Thus, we do not need to consider inverse data augmentation or alignment. In our experiments, we used horizontal flipping and translation as image augmentations. However, some augmentation methods, such as “Simple Copy–Paste”, can completely change the instance class in images, which conflicts with contrastive learning and is not suitable for overuse in our method.

3.3. Loss Transform

Given a loss function L, we propose LossTransform by incorporating an additional regularization component. For a mini-batch of predictions $\{p_1,p_2,\dots ,p_n\}$ with corresponding targets $\{y_1,y_2,\dots ,y_n\}$, the model parameters are optimized by minimizing the following objective function:

$$\mathcal{T}\left(L\right)={\displaystyle \frac{1}{n}}\sum _{i=1}^n\left[L(y_i;p_i)+{\displaystyle \frac{1}{m}}\sum _{j=1}^m\rho D_{\mathrm{KL}}(\overline{p_i}\left|\right|{\psi }_j^{-1}\left(p_i\right))\right],$$

(12)

where we chose $\rho =1$ as the default setting because it provides a stable trade-off across different tasks and datasets, and ${\overline{p}}_i$ represents the average prediction over all instances positively paired with ${\psi }^{-1}\left(p_i\right)$. $T\left(L\right)$ consists of two components: the traditional mini-batch loss on the left, and the distillation loss on the right, which is computed as the average of m Kullback–Leibler (KL) divergences.

4. Results

In this section, we start by evaluating the effectiveness of LossTransform in enhancing the generalizability of classifiers. Our experiments are primarily conducted on the CIFAR-10/100 and ImageNet-1K datasets. Initially, we investigate the factors that affect the performance of LossTransform. Subsequently, we demonstrate that models trained with LossTransform yield more stable predictions. Additionally, we employ LossTransform methods on various datasets and classifiers for comparison with baselines. We then explore the applications of LossTransform in the fine-tuning process.

Next, we analyze the memory limitations of the LossTransform method. Due to the increased computational demand of duplicating sample replicas for data augmentation, training large models on datasets with large-scale images requires additional computational resources. To address this limitation, LossTransform-v1 is introduced, employing subsampling of Repeated Augment Sampling, thereby not increasing the computational resources required for model training. Consequently, it offers a feasible solution to the limitations of LossTransform. Experiments on ImageNet and Microsoft COCO reveal that fine-tuning pre-trained models with LossTransform-v1 can yield better evaluations than state-of-the-art baselines.

In our experiments, we adhere to the original training recipes of all benchmark models, including their optimizers, batch sizes, epochs, and learning rate schedules.

4.1. Training Benefits from LossTransform

4.1.1. Suitability for Complex Augmentation

Here, we verify that the LossTransform of Categorical Cross Entropy (CCE) works well across all data augmentation schemes. Table 1 presents a performance comparison between LossTransform applied to Categorical Cross Entropy $\mathcal{T}\left(\mathrm{CCE}\right)$ and the traditional CCE. For each combination of data augmentation schemes, we follow the official training recipe in the ResNet paper and apply CCE and $\mathcal{T}\left(\mathrm{CCE}\right)$ for model training. All experiments are repeated five times with result variation under 0.3%.

Table 1. Test error rates (%) of plain/LossTransform-trained models on CIFAR-10. The affine transformation [26] and horizontal flipping schemes are used by default. Cutout [27], CutMix [28], and mixup [29] are optional. The number of repetitions for Repeated Augmentation Sampling is set to 32. To emphasize general patterns, the best test error rates are highlighted in bold.

Model	Cutout	CutMix	Mixup	CCE	$\mathcal{T}\left(\mathbf{CCE}\right)$
	✗	✗	✗	7.07	6.46
ResNet-20	✓	✗	✗	7.23	6.70
no. of params = 0.27 M	✗	✓	✗	7.50	6.13
	✗	✓	✓	7.70	6.27
	✗	✗	✗	6.32	5.73
ResNet-40	✓	✗	✗	6.22	5.58
no. of params = 0.66 M	✗	✓	✗	5.99	4.47
	✗	✓	✓	7.70	4.39
	✗	✗	✗	6.10	5.64
ResNet-110	✓	✗	✗	5.71	4.90
no. of params = 1.7 M	✗	✓	✗	5.07	4.28
	✗	✓	✓	5.76	4.02

The last column of Table 1 demonstrates that $\mathcal{T}\left(\mathrm{CCE}\right)$ is efficient for all hybrid data augmentation schemes. Empirically, while the traditional CCE-trained model shows improvement with a single data augmentation scheme, it fails under hybrid schemes due to the ambiguities introduced by mixed data augmentations. $\mathcal{T}\left(\mathrm{CCE}\right)$ overcomes this limitation to work for all data augmentation combinations and improve the model performance.

4.1.2. Faster Convergence

A loss function that converges to lower local minima reflects enhanced model performance, while faster convergence reduces computational resources, improving training efficiency. The LossTransform converts a loss function into a better form, thereby increasing the training speed and improving better model performance. Figure 2 shows a typical example of the CCE loss function, where models gain additional performance from the faster training on $\mathcal{T}\left(\mathrm{CCE}\right)$.

Figure 2. Test error of ResNet-20 and DenseNet-40 on CIFAR-10 dataset. The model trained with LossTransform effectively gains lower test errors and accelerates the training process.

Based on CCE and the official training recipe, the top-1 validation accuracy of DenseNet has stopped improving after 300 epochs. However, our loss function $\mathcal{T}\left(\mathrm{CCE}\right)$ converges faster with 120 epochs. Figure 2 (left) illustrates the test error (%) of CCE-trained ResNet and $\mathcal{T}\left(\mathrm{CCE}\right)$-trained ResNet, respectively, where $\mathcal{T}\left(\mathrm{CCE}\right)$ saves 25% training iteration steps. Figure 2 (right) shows the same performance on DenseNet where $\mathcal{T}\left(\mathrm{CCE}\right)$ saves 60% training iteration steps.

4.1.3. Stable Predictions

We demonstrate that classifiers trained using LossTransform have a significant change in predictive behaviour. Figure 3 shows the difference between a standard-trained model and a LossTransform-trained model.

Figure 3. The predicted probabilities of eight augmented images. The left $8\times 10$ matrix corresponds to the predicted probabilities of a ResNet-20 model trained with CCE loss. The right $8\times 10$ matrix corresponds to the predicted probabilities of a LossTransform-trained ResNet-20 model. We color the probability background with its value; the higher the value, the deeper the square.

The standard-trained model tends to assign inconsistent predicted probabilities across different augmented images; therefore, the predicted results are easily affected by data augmentation (see Figure 3, left). On the contrary, when we apply LossTransform, the model outputs almost identical prediction probabilities to different augmented images (see Figure 3, right). To the image classification task, we now introduce a concrete variance metric to measure the stability of predictions:

$$\mathrm{Var}\left(\pi \right)={\displaystyle \frac{1}{m-1}}{∥\sum {\pi }_i-{\displaystyle \frac{1}{m}}\sum _{j=1}^m{\pi }_j∥}^2$$

(13)

where $\{{\pi }_1,\cdots ,{\pi }_m\}$ are the prediction probability vectors of an image under m different augmentations. Figure 3 shows that the model trained with LossTransform has stable output and is minimally influenced by noise and data augmentation.

4.1.4. The Number of Repetitions for REPEATED Augmentation

To evaluate the impact of the number of repetitions m on model training, four distinct data augmentation techniques are used, including affine transformation, horizontal flipping, CutMix, and mixup, with different values of m during training. For other training settings in our experiments, we follow the configurations outlined in the ResNet and DenseNet papers.

Figure 4 illustrates that the test accuracy reward of models is positively correlated with m. The model trained with $\mathcal{T}\left(\mathrm{CCE}\right)$ tends to yield consistent improvement in accuracy with the growing m. Another factor that affects performance is the network complexity. The model with fewer parameters gains higher test accuracy improvement when training with the $\mathcal{T}\left(\mathrm{CCE}\right)$ loss function. In Appendix B, we also present ablation and sensitivity analyses for both parameters $\rho$ and m.

Figure 4. Left: All networks tend to gain more performance improvement with the growing number of repetitions m. Right: Small models gain more improvement in CIFAR-10@Acc1 and CIFAR-100@Acc5 tests.

4.1.5. Comparison with Other Loss Functions

Table 2 presents the comparison between the LossTransform and other methods when training models from scratch. We achieve better accuracies on all mini-size ResNet and DenseNet models without using extra training data or pre-trained models.

Table 2. Test error rates on CIFAR-10 and CIFAR-100 datasets. Comparisons with default CCE (baseline) and Supervised Contrastive Learning (SupCon). The best rates are highlighted in bold.

Model	Loss Function	CIFAR-10	CIFAR-100
ResNet-20 no. of params = 0.27 M	CCE	7.70%	33.60%
	SupCon	6.55%	31.10%
	$\mathcal{T}\left(\mathrm{CCE}\right)$	6.27%	28.59%
ResNet-40 no. of params = 0.66 M	CCE	6.68%	29.80%
	SupCon	5.23%	26.48%
	$\mathcal{T}\left(\mathrm{CCE}\right)$	4.39%	25.34%
ResNet-110 no. of params = 1.7 M	CCE	5.76%	28.45%
	SupCon	4.56%	25.14%
	$\mathcal{T}\left(\mathrm{CCE}\right)$	4.02%	24.37%
DenseNet-40 no. of params = 1.0 M	CCE	5.67%	27.37%
	SupCon	4.52%	24.90%
	$\mathcal{T}\left(\mathrm{CCE}\right)$	4.35%	23.55%
ResNet-50 no. of params = 25.6 M	CCE	5.0%	29.7%
	SimCLR	6.4%	29.3%
	Max-Margin	7.6%	29.5%
	SupCon	4.0%	23.5%

We also list in Table 2 the results of the large-scale ImageNet model (ResNet-50 with 25.6 million parameters) on the CIFAR dataset [2], the experiments of comparing cross-entropy training, SimCLR [1], maxi-margin classifiers [30], and SupCon [2]. In comparison, we achieve better results on smaller models, and our loss function only requires 120 epochs for training, while other methods are much more complicated (≈1000 epochs).

LossTransform is also suitable for fine-tuning. We conduct experiments using pre-trained MobileNet-V2 [31], NASNetMobile [32], and EfficientNet-B0 [33] models on the ImageNet-1K dataset. In these experiments, we still use the aforementioned simple data augmentation schemes for fair comparison. All models are trained end-to-end by SGD with a momentum of 0.9, a constant learning rate of 0.001, and 30 epochs. We use a batch size of 128 and set m for up to 4.

Table 3 shows the fine-tuning results of $\mathcal{T}\left(CCE\right)$ for deeper networks. The top-1 and top-5 accuracy refer to the model’s performance on the ImageNet-1K validation dataset. Remarkably, all the pre-trained models are downloaded from the TensorFlow platform. The real validation accuracy (without using any TTA and ensemble methods) of these pre-trained models is much lower. The validation accuracies (MobileNet-V2: 71.3%; NASNetMobile: 74.4%; EfficientNet-B0: 77.1%) shown on the TensorFlow webpage and in corresponding papers are the results of ensemble and multi-crop models [34]) (MobileNetV2: 67.59%; NASNetMobile: 72.11%; EfficientNet-B0: 76.22%).

Table 3. Comparison among the pre-trained models and fine-tuned models on the ImageNet-1K validation set.

Model	Loss Function	Acc@1	Acc@5
MobileNet-V2	CCE	70.81%	89.93%
	$\mathcal{T}\left(\mathrm{CCE}\right)$	71.26%	90.48%
NASNetMobile	CCE	72.35%	90.87%
	$\mathcal{T}\left(\mathrm{CCE}\right)$	73.14%	91.54%
EfficientNet-B0	CCE	74.57%	92.13%
	$\mathcal{T}\left(\mathrm{CCE}\right)$	77.09%	93.61%

4.1.6. Object Detection and Instance Segmentation

Unlike image classification, object detection and instance segmentation are advanced tasks in computer vision. More specifically, object detection is a hybrid task of regression and classification that requires detecting objects of interest in unknown numbers. Instance segmentation needs to classify each pixel in the image to provide more fine analysis results. To date, most learning algorithms are still evaluated on benchmarks of classification tasks, so building a generalized algorithm that can traverse AI fields is an important challenge.

We investigate the application of LossTransform on PennFudanPed object detection and instance segmentation task. Table 4 shows the evaluation of models on the test set. We get significant improvements by using LossTransform applied to the multi-task loss function $\mathcal{T}\left(l_{\mathrm{multi}}\right)$, where $l_{\mathrm{multi}}$ is the plain multi-task loss function in Mask R-CNN. In this table, every pair of experiments uses the same training configurations, in which each model is trained by two different loss functions, respectively. So the precision gain of models is entirely due to the LossTransform method.

Table 4. Comparison with the multi-task loss function of Mask R-CNN on PennFudanPed object detection and instance segmentation. Horizontal flipping, random cropping, and simple copy–paste data augmentations are used, and all models are trained from scratch. Mask R-${\mathrm{CNN}}^{*}$ is initialized with a MS COCO pre-trained model and fine-tuned on the new dataset.

Backbone	Loss	Detection (%)			Segmentation (%)
		AP	AP50	AP75	AP	AP50	AP75
ConvNeXt-tiny	$l_{\mathrm{multi}}$	65.4	96.7	82.7	60.5	91.5	79.2
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	65.6	96.9	83.6	63.5	94.0	81.8
DensNet121	$l_{\mathrm{multi}}$	64.6	96.7	74.9	62.8	96.7	74.8
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	67.1	99.8	82.7	68.7	99.8	88.3
EfficientNetB4	$l_{\mathrm{multi}}$	73.8	100	91.3	66.8	96.8	88.7
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	74.4	100	91.4	69.0	100	85.6
GoogLeNet	$l_{\mathrm{multi}}$	59.6	98.9	74.0	60.1	98.9	75.8
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	59.4	96.6	68.7	63.1	96.1	81.6
InceptionV3	$l_{\mathrm{multi}}$	71.0	99.7	87.0	64.6	99.5	82.2
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	75.0	99.8	91.5	66.9	99.8	88.8
MNASNet1.0	$l_{\mathrm{multi}}$	61.7	94.1	78.3	56.0	94.1	65.7
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	65.5	96.5	81.1	61.0	93.7	83.0
ResNet50	$l_{\mathrm{multi}}$	66.1	99.5	78.4	64.6	99.4	80.4
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	70.7	99.9	88.6	70.1	99.9	91.8
ShuffleNetV2	$l_{\mathrm{multi}}$	70.5	100	83.9	62.7	96.8	81.8
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	72.0	99.3	90.9	64.0	96.5	85.6
WideResNet50	$l_{\mathrm{multi}}$	69.0	100	84.5	70.8	100	91.2
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	74.2	100	90.9	71.3	100	92.2
Mask R-${\mathrm{CNN}}^{*}$	$l_{\mathrm{multi}}$	86.9	100	96.8	80.6	100	96.8
	$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	88.2	100	96.7	81.9	100	96.7

In the first $9\times 2$ experiments, we replace the feature pyramid network (as of v0.14, TorchVision offers the model architectures for ImageNet, and we extract their backbone with a feature pyramid network (FPN) on top of Mask R-CNN) [35] of Mask R-CNN architecture with different structures and retrain the whole model from scratch. For DenseNet121- [14], GoogLeNet- [36], ResNet-50-, and WideResNet50-based [37] FPNs, we set anchors to ((32,), (64,), (128,), (256,), (512,)) which means there are five levels/groups with one anchor size. For ConvNeXt-tiny- [38], EfficientNetB4- [33], InceptionV3- [39], MNASNet1_0- [40], and ShuffleNetV2-based [41] FPNs, we set anchors to ((32, 64, 128, 256, 512), (32, 64, 128, 256, 512)). These models are trained with a mini-batch size of 4 and a total of 18 epochs, with the learning rate divided by factor 10 at epochs 12 and 16. For DensNet121, EfficientNetB4, GoogLeNet, ResNet50, ShuffleNetV2, and WideResNet50, we start training with an initial learning rate of 0.01. Other models start with 0.0025.

In the last two experiments, we freeze the FPN and region proposal network [5] of the pre-trained Mask R-CNN model and then fine-tune it on the new dataset. The learning rate starts at 0.005.

We further investigate the PASCAL VOC 2012 dataset. In order to train the Mask R-CNN with LossTransfrom on this dataset, we need to convert the VOC to a dataset that can conduct the instance segmentation task. Here, we preprocess the dataset by using the algorithm [11,42] to split distinct individual instances from multiple objects of the same class. Then, we freeze the ResNet50-FPN and the region proposal network (RPN) of Mask-RCNN and use transfer learning to retrain the RoI heads. Finally, we achieve 73.5% mAP on the object detection task by training the model with $\mathcal{T}\left(l_{\mathrm{multi}}\right)$, which gains an extra +0.8% mAP than using $l_{\mathrm{multi}}$.

4.2. Fine-Tuning Benefits from LossTransform-v1

From Figure 1, it can be seen that the reformulated loss function in LossTransform consumes multiple times more memory during computation. For example, a mini-batch of images $\{x_1,x_2\}$, after repetition and data augmentation, results in the actual training input $\{{\psi }_1\left(x_1\right),{\psi }_2\left(x_1\right),{\psi }_3\left(x_1\right),$${\psi }_4\left(x_2\right),{\psi }_5\left(x_2\right),{\psi }_6\left(x_2\right)\}$, leading to a data throughput during optimization that is three times greater than in the standard training phase. This implies that LossTransform can effectively train large models on large-scale image datasets only if sufficient GPU memory is available.

To address this memory limitation, we apply RASampler in LossTransform. RASampler includes a shuffling and subsampling process. For instance, after sequentially applying shuffling, subsampling, and data augmentation, the actual mini-batch images become $\{{\psi }_1\left(x_1\right),{\psi }_2\left(x_2\right),{\psi }_3\left(x_1\right)\}$, ensuring that the number of examples the model learns from in each iteration remains unchanged.

We investigate the application of LossTransform-v1 on the ImageNet-1K classification (see Table 5) and MS COCO object detection (see Table 6) tasks. In our fine-tuning experiments, TorchVision provides multi-version pre-trained weights for each model architecture [43]. For fair ablation studies, we use the “V1 pre-trained weights” as the default since the corresponding training recipes for “V1 pre-trained weights” are not specifically tailored to each model architecture. We only keep the learning rate as the value from the last epoch in recipes. We empirically find that, when using LossTransform-v1, the model typically achieves performance improvement within 20 training epochs. Our experiments show that LossTransform is effective even for complex models like ViT [44].

Table 5. Comparison among the pre-trained and fine-tuned models on the ImageNet-1K validation set.

Backbone	Baseline (%)		LossTrans (%)
	Acc@1	Acc@5	Acc@1	Acc@5
ResNet-50	76.130	92.862	76.206	93.138
ResNext-50	77.618	93.698	77.853	93.923
MobileNet-V2	71.878	90.286	72.020	90.310
EfficientNet-B0	77.692	93.532	77.817	93.674
ViT-B16	81.072	95.318	82.019	95.326

Table 6. Comparison among the pre-trained and fine-tuned models on MS COCO object detection task. Horizontal flipping data augmentation is used, and all model heads are trained from scratch.

Model	No. of Params	Loss	Box MAP
Faster R-CNN	41.8M	$l_{\mathrm{multi}}$	37.0%
ResNet-50 FPN		$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	37.1%
Faster R-CNN	19.4M	$l_{\mathrm{multi}}$	32.8%
MobileNetV3-L FPN		$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	33.2%
RetinaNet	34.0M	$l_{\mathrm{multi}}$	36.4%
ResNet-50 FPN		$\mathcal{T}\left(l_{\mathrm{multi}}\right)$	36.6%

5. Discussion

It is important to acknowledge the limitations of our method. Currently, we focus exclusively on image classification and object detection tasks. However, other recognition or generation tasks, such as text classification and causal language modeling, are also fundamental classification problems. This suggests that our method may be applicable beyond the vision domain, though empirical validation is needed in future work.

In addition, our approach duplicates input data multiple times within each mini-batch to perform diverse data augmentations; while this strategy improves training dynamics, as shown in Section 4, it inevitably increases GPU memory consumption. Nonetheless, we observe that our method accelerates model convergence compared to standard training, which may offset the additional memory overhead over the course of training.

Recent advances in contrastive learning, such as MoCo v3 [45] and DINO [46], have demonstrated strong performance in classification and object detection tasks. Although MoCo v3 and DINO achieve better results in some scenarios, MoCo v3 requires dual encoders for contrastive loss, and DINO relies on an additional teacher network for distillation. Both approaches introduce additional trainable parameters and increase model size. In contrast, LossTransform is a novel loss function paradigm that requires no architectural modifications and can be applied to any model using cross-entropy or multi-loss frameworks.

Finally, our method has not yet been applied to causal language modeling or text classification: (1) LossTransform relies on multiple parallel stochastic augmentations of the same input, which conflicts with causal language modeling. Despite both image classification and causal language modeling using cross-entropy as a loss function, for supervised fine-tuning (SFT) [47] in large language models (e.g., Gemma [48], Qwen [49], DeepSeek [50]), data augmentation is usually prohibited. (2) Perspectively, our method remains applicable to downstream text classification tasks (and other discriminative language tasks), where data augmentation strategies (e.g., back-translation, synonym replacement) are commonly used. In this scenario, the output port of the text classification model is identical to the image classification model, both using cross-entropy loss for optimization, with data augmentation allowed.

6. Conclusions

In this paper, we propose the LossTransform method for training neural networks. This method represents a more generalized approach to contrastive learning, which can be easily implemented through the reformulation of the loss function. It has proved to be effective across various data augmentation schemes, recognition tasks, and neural network architectures. The core principle of LossTransform is to minimize the distances between positive samples or pull together positive samples to their center.

LossTransform does not require the modification of data pipelines, model architectures, and training recipes. Therefore, it will not increase the computational resources requirement for image classification/object detection models. Any model that can be deployed on resource-constrained environments (smartphones or edge devices) can also benefit from LossTransform after training/fine-tuning on servers. This method is user-friendly with a plug-and-play implementation, as it upgrades only the loss function in practical scenarios. It is applicable for either training from scratch and fine-tuning.

For model training, LossTransform contributes to improving the stability of model inference, boosting the convergence speed, and enhancing the efficiency of utilizing data augmentation. LossTransform-v1 addresses the shortcomings of LossTransform regarding memory constraints. In the case of challenging recognition tasks, LossTransform-v1 can continue to yield better evaluation than state-of-the-art baselines.