Exploring the Impact of Temperature Scaling in Softmax for Classification and Adversarial Robustness

TL;DR

This work investigates the temperature parameter in softmax for image classification, combining gradient analysis with extensive experiments on CNNs and transformers across CIFAR-10/100 and Tiny-ImageNet. It shows that moderate $\tau$ can enhance overall accuracy, while elevated $\tau$ reshapes gradient flow to promote balanced learning and unexpectedly boosts robustness against common corruptions and untargeted PGD attacks, with potential benefits for adversarial training. The findings offer a simple yet impactful hyperparameter for improving performance and security in deep learning, while acknowledging limitations in selecting an optimal $\tau$ automatically.

Abstract

The softmax function is a fundamental component in deep learning. This study delves into the often-overlooked parameter within the softmax function, known as "temperature," providing novel insights into the practical and theoretical aspects of temperature scaling for image classification. Our empirical studies, adopting convolutional neural networks and transformers on multiple benchmark datasets, reveal that moderate temperatures generally introduce better overall performance. Through extensive experiments and rigorous theoretical analysis, we explore the role of temperature scaling in model training and unveil that temperature not only influences learning step size but also shapes the model's optimization direction. Moreover, for the first time, we discover a surprising benefit of elevated temperatures: enhanced model robustness against common corruption, natural perturbation, and non-targeted adversarial attacks like Projected Gradient Descent. We extend our discoveries to adversarial training, demonstrating that, compared to the standard softmax function with the default temperature value, higher temperatures have the potential to enhance adversarial training. The insights of this work open new avenues for improving model performance and security in deep learning applications.

Figures (7)

• Figure 1: Demonstration of the model optimization direction with different temperatures. $f(x)$ is the latent code of a data sample from category 3. Since $f(x)$ is close to the negative class prototype $w_1$, the CE loss with respect to the encoder $f$ yields a large gradient toward the groundtruth $w_3$. However, with different temperature factors, the gradients associated with the negative classes are different: low temperature makes the update more biased by the hard class (a), while an elevated temperature leads to more equalized gradients (b).
• Figure 2: T-SNE tsne visualization of the CIFAR10 sample distribution after the ResNet50 encoder with different temperatures.
• Figure 3: T-SNE tsne visualization of the CIFAR10 sample distribution after the VIT encoder with different temperatures.
• Figure 4: Test error number during training. The red line represents $\tau=0.5$, the green line represents $\tau=1$, and the orange line represents $\tau=50$. The model used is Resnet50 and is tested on CIFAR10. SGD optimizer is used during training with the learning rate set to 0.1 (a) and 0.01 (b). The shade areas consist of 6 total runs with different random seeds. The solid lines indicate the mean value across all runs.
• Figure 5: The logit changes before and after PGD20 attack. The blue lines stand for the logits of the samples before PGD attack, and the orange lines stand for the logits of the samples after PGD attack.