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

Given a set of real numbers, $X=\{x_{1},..,x_{N}\}$, the generalized softmax function can be used to normalize $X$ into a probability distribution.

where $\mathbb{S}$ represents the softmax function and $\tau$ is the temperature scaling factor. The temperature $\tau$ controls the smoothness (softness) of the probability it produces. Specifically, when $\tau\rightarrow\infty$, the output tends toward a uniform distribution; while when $\tau=0$, the softmax function assigns a probability of 1 to the element with the highest value and a probability of 0 to the rest. The standard (unit) softmax function, with $\tau=1$, is widely used in conventional classification tasks.

We consider multi-category classification in this study, where paired training data $\{\mathscr{X},\mathscr{Y}\}=\{(x,y)|x\in\mathbb{R}^{H\times L\times N},y\in\mathbb{R}^{1\times M}\}$ are drawn from a data distribution $\mathcal{D}$. Here, $H,L,N$ are the dimension of a sample $x$, $M$ is the number of categories, and $y$ is a one-hot vector indicating the class of the input $x$. A classifier, $\mathcal{C}:\mathscr{X}\xrightarrow{}\mathscr{Y}$, is a function predicting the label $y$ for a given data $x$. That is $C(x)=y$. In the canonical classification setting, a neural network classifier, $\mathcal{C}=(f,W)$, is usually composed of a feature extractor $f$ parameterized by $\theta$ and a weight matrix $W$. $f$ is a function mapping the input $x$ to a real-valued vector $f(x)$ in the model’s penultimate layer and $W=(w_{1},...,w_{M})$ represents the coefficients of the last linear layer before the softmax layer. So the likelihood probability of data $x$ corresponding to the $M$ categories can be formulated as

Note that each vector $w_{i}$ in matrix $W$ can be considered as the prototype of class $i$ and the production $W^{T}f(x)$ in Eqn. 2 quantifies the similarity between the feature $f(x)$ and different class-prototypes.

During training, the model $\mathcal{C}=(f,W)$ is optimized to minimize a specific loss, usually a Cross-Entropy (CE) loss.

Though $\tau=1$ is the default setting in classification tasks, we preserve $\tau$ in the Eqn.s to facilitate theoretical analysis.

We conduct image classification on multiple benchmarks (i.e. CIFAR10, CIFAR100, and Tiny-ImageNet) and their extended Common Corruptions and Perturbations sets (i.e. CIFAR10-C, CIFAR100-C, and Tiny-ImageNet-C with corruption strentgh being 3) to investigate the impact of temperature scaling. In addition, we also evaluate the model’s robustness against adversarial attacks such as PDG20 [18] and C&W [2]. Both attacks are bounded by the $l_{\infty}$ box with the same maximum perturbation $\epsilon=8/255$.

To get a comprehensive evaluation, we set $\tau\in\{0.1,0.5,1,10,30,50,70,100\}$. Unless stated otherwise, we takes ResNet50 and VIT-small-patch16-224 as the CNN and transformer backbones, respectively. The ResNet50 is trained from scratch, with SGD optimizer and learning rate setting to 0.1. We also utilize the Cosine Annealing scheduler to better train the model. The transformer is pretrained on ImageNet-21K and finetuned on the target dataset using Adam optimizer. All experiments run on one RTX3090.

To clarify, the temperature scaling only involves in model training in this study, but not model evaluation and attacks. All empirical evaluation and adversarial sample generation by PGD and C&W are based on the standard cross entropy, i.e. $\tau=1$. Thus, attack gradients are not attenuated, reflecting model’s true sensitivity to data perturbation.

The quantitative results on CNN and Transformer are summarized in Table I and Table II, respectively. For the CNN model, ResNet50, training from scratch, the standard accuracy increases with the temperature increase. Furthermore, CNN models trained at elevated temperatures show more robustness against naturally corrected images. We believe that such improvements are majorly attributed to better model optimization with leveraged temperature. For the transformer finetuned on the target set, the standard accuracy and robustness against natural corruptions and perturbations is quite stable. We hypothesize that such stable performance is due to the fact that ViT has already been pre-trained on ImageNet and has reached a relatively high-quality state. Additionally, we observed that the model’s adversarial robustness gradually improves with increasing temperature.

Clustering is a crucial metric when measuring how an encoder performs. In classification, a good encoder should be able to gather samples from the same class while separating clusters of different classes. Fig. 2 and Fig. 3 present 2D TSNE visualization of the CIFAR10 sample distribution by ResNet50 and transformer. We observe a similar trend: low temperatures lead to more mixed clusters, while models trained with elevated temperatures have better cluster effects. These empirical observations also explain the improved classification performance on clean and non-adversarial perturbations, as well as stronger adversarial robustness, with high temperature in Table I and Table II.

We then conduct experiments observing the training process when applying different temperatures to the model. We validate the model on the test set every epoch and record the error probability. As our results shown in Fig. 4(a), we can clearly observe that not only does the training convergence speed increase as the temperature goes up, but models trained with higher temperatures also tend to converge to lower points, leading to better final performance. In fact, when we further decrease the temperature to around 0.1, the model would have a substantial risk of not converging at all. While this might appear contrary to the common understanding that focusing on hard classes will generally benefit the model, a more nuanced explanation is provided by delving further into the gradient analysis provided in Section 4.

From Eqn. 4, 5, 6, we observe that if the logit of the target class is not the largest, its gradient will increase dramatically with low temperatures. This is potentially bad for models being known to converge inefficiently under large learning rates. One straightforward solution would be lowering the learning rate as shown in Fig. 4. While the training converging speeds are closer, the run with a higher temperature can still reach a better performance. Furthermore, regardless of the increase in overall training converging speed for $\tau=0.5$ and $\tau=1$ runs when lowering the learning rate, the final performances for all three runs actually get worse than runs with 0.1 learning rate. Therefore, this phenomenon cannot be attributed solely to a high learning rate. However, if we shift our perspective to the overall direction for optimization as done in Eqn. 9, it becomes clear that during the early stage of training, the encoder $f$ has not converged to an ideal point, leading to sub-optimal values produced for certain update directions. If this happens to be the direction of the target class and the error-prone class which models with small temperatures tend to focus on, the model training can be impacted harmfully. In the meantime, high temperatures equalize the weight given to all the classes and ensure the update is not terribly wrong even if a few ${\partial L_{ce}(x)}/{\partial w_{j}}$ are in the wrong direction. Upon reaching this conclusion, we are surprised to find that this reasoning and our empirical observations align perfectly with the curriculum learning philosophy, that starting from hard samples may harm model optimization and learning outcomes.

Table I and Table II show that models trained with elevated temperatures have strong adversarial robustness. TSNE plots in Fig. 2 and Fig. 3 also support this observation. This prompts questions regarding the mechanism behind the gained robustness. In this section, our focus is on investigating the model’s behavior under adversarial attacks and understanding why the model demonstrates such robustness.

Gradient analysis for adversarial generation. In order to discern the source of model robustness, we follow the work in [9] and study the gradient of the classification loss with respect to the input to analyze the direction of the PGD attack, which can be written as

As illustrated above, given a well-trained model, for most inputs where $\mathbb{S}(w^{T}_{i}\cdot f(x))\approx 1$, the gradient does not have a noticeable portion in target class $w_{i}$ on the early stage of the attack. This implies that rather than directly ’stepping away’ from the target class, the attack will initially focus on approaching other class prototypes. Moreover, the second term, $\sum_{j\neq i}w^{T}_{j}\cdot\mathbb{S}(w^{T}_{j}\cdot f(x))$, indicates that all the other directions are weighted by their according probabilities. Therefore, untargeted attacks are actually targeted toward the error-prone class, which most commonly is the largest probability class other than the target class. However, if a model lacks an error-prone class given an input, all $w_{k}$ will be weighted equally. Consequently, the gradient would point toward all negative class prototypes, making it exceptionally challenging to determine the optimal direction. We noticed that such a scenario occurs when a model is trained with a small $\tau$. Then let’s focus on the gradient update strength. For a data sample $x$ is classified correctly, $\mathbb{S}(w^{T}_{j}\cdot f(x))$ would be small when the model training temperature $\tau$ increases. That is, when a model is trained with high temperatures, not only the gradient direction to generate adversarial samples is not clear, but the gradient strength is also small. Both factors contribute to the robustness of the model when optimized with elevated temperatures.

Raw Logit Analysis. With the insight from the gradient analysis on adversarial attack, we then turn to observe the logit output around the adversarial attack, as shown in Fig. 5. Each bar represents the logit value for each class, blue bars stand for the logit outputs of clean samples and orange bars are the logit outputs from adversarial samples. Models share similar characteristics in low temperatures, Fig. 5(a,b), with the logit of the target class going down while the logit of the error-prone class going up. However, for models trained with large temperatures, Fig. 5(c,d), two logits are nearly identical with a minimal amount of changes. This contrasts the robustness gains during adversarial training, where the model learns the pattern of the adversarial noise.

Class Prototypes Analysis. To further analyze the model behavior, we investigate the relation between the encoded feature, $f(x)$, and each class prototype, $w_{j}$. Here, we observe the Euclidean distance and cosine similarity. Fig. 6 shows Euclidean distance and cosine similarity between one sample and all class prototypes. It is evident that as the training temperature goes up, the feature $f(x)$ tends to have an identical distance to all negative class prototypes. This indicates the model trained with high temperature is less likely to have an error-prone class, which is essential for untargeted attacks as we discuss above.

Furthermore, to illustrate that the phenomenon shown in Fig. 6 is not limited to one or a few samples, we calculate the variance of Euclidean distance and cosine similarity of all negative class prototypes across all samples in CIFAR10 test set. Note that as illustrated in Fig. 6, different models have very different ranges for Euclidean distance between encoded feature and class prototypes. Therefore, we map the value of different models into the same range to make a more direct comparison. Box plots are drawn in Fig. 7 showing the overall variance results with each box being a model trained with a different temperature. We can observe a clear trend that when the temperature rises, the variance for both Euclidean distance and cosine similarity drops indicating the encoded sample, $f(x)$, has a more similar distance to all negative class prototypes. One might notice an increase in variance when the temperature reaches some threshold. We label them as extreme temperatures, which are so large that they can adversely affect the model’s convergence.

Despite the model trained with high temperatures showing superb robustness against untargeted PGD attack due to its nature attribute that discovers the weakness of PGD attack, it does not hold robustness against targeted attacks. The reason behind this is straightforward. In targeted attacks, Eqn. 10 no longer holds, and the gradient is not obligated to move towards all negative class prototypes with a weighted step size. Therefore, with the only source of the model robustness gained eliminated, it is naturally vulnerable to targeted attacks.

Remark: Even though many attacks claim themselves to be untargeted attacks, they actually optimize toward one self-selected target, which we do not consider untargeted attacks under this setting. One popular example is the Difference of Logits Ratio(DLR) attack proposed by  [4]. Regardless of its ability to rescale the logit,

shows that the DLR loss automatically selects the class holding the largest logit other than the target class as the attack target. Therefore, during optimization, it does not need to optimize toward all negative class prototypes. A similar example also includes FAB attack [3].

Given that our temperature control method is used inside the Cross-Entropy Loss, it is possible to apply this method in adversarial training. Here, we do preliminary experiments on the adversarial training baseline proposed by  [18] for the simplicity of its loss function. We add temperature control inside vanilla loss term forming

where $F$ is a combination of encoder and class prototypes.

Our preliminary results are listed in Table III. We can clearly observe that model robustness increases as the temperature increases with a slight trade-off with clean accuracy, which confirms the possibility of combining the temperature control method with adversarial training. While further extension to other adversarial training methods is possible, it remains a complex problem for most adversarial training involves complex loss functions that may introduce terms other than the Cross-Entropy function. Also, balancing the vanilla loss term and adversarial loss term largely relies on empirical experiments. Therefore, further exploration of fitting this into other adversarial training methods falls beyond the scope of this paper.