[论文解读] Temperature check: theory and practice for training models with softmax-cross-entropy losses

Atish Agarwala, Jeffrey Pennington+2 more|arXiv (Cornell University)|Oct 14, 2020

Domain Adaptation and Few-Shot Learning参考文献 32被引用 13

一句话总结

本文为使用softmax交叉熵损失训练的模型的早期学习动态建立了一套理论，表明逆温度 $β = 1/T$ 对训练速度和泛化能力具有关键影响，且该影响与初始logit幅度无关。在CIFAR10、ImageNet和IMDB上的实证结果表明，将 $\beta$ 调节在 $[10^{-2}, 10^1]$ 范围内可显著提升性能，且最优 $\beta$ 值通常远非1，表明温度是影响模型性能的关键超参数。

ABSTRACT

The softmax function combined with a cross-entropy loss is a principled approach to modeling probability distributions that has become ubiquitous in deep learning. The softmax function is defined by a lone hyperparameter, the temperature, that is commonly set to one or regarded as a way to tune model confidence after training; however, less is known about how the temperature impacts training dynamics or generalization performance. In this work we develop a theory of early learning for models trained with softmax-cross-entropy loss and show that the learning dynamics depend crucially on the inverse-temperature $β$ as well as the magnitude of the logits at initialization, $||β{\bf z}||_{2}$. We follow up these analytic results with a large-scale empirical study of a variety of model architectures trained on CIFAR10, ImageNet, and IMDB sentiment analysis. We find that generalization performance depends strongly on the temperature, but only weakly on the initial logit magnitude. We provide evidence that the dependence of generalization on $β$ is not due to changes in model confidence, but is a dynamical phenomenon. It follows that the addition of $β$ as a tunable hyperparameter is key to maximizing model performance. Although we find the optimal $β$ to be sensitive to the architecture, our results suggest that tuning $β$ over the range $10^{-2}$ to $10^1$ improves performance over all architectures studied. We find that smaller $β$ may lead to better peak performance at the cost of learning stability.

研究动机与目标

• 理解softmax温度 $T$（或其逆 $\beta$）如何影响使用交叉熵损失训练的深度神经网络中的学习动态与泛化能力。
• 将温度 $\beta$ 的影响与初始logit幅度 $\|\mathbf{Z}^0\|_2$ 对训练性能的影响分离开来。
• 探究在非单位 $\beta$ 下泛化能力的提升是否源于模型置信度的变化，或源于早期训练过程中的动力学效应。
• 在多样化架构与数据集上，实证验证理论预测的学习时间尺度，并识别最优 $\beta$ 范围。
• 确立 $\beta$ 为可调超参数，其调节可显著提升模型性能，超越标准的 $T=1$ 设置。

提出的方法

• 基于初始化附近的线性化，对早期学习动态进行理论分析，通过 $\tilde{\eta} = \eta \beta^2$ 对梯度进行缩放，以解耦 $\beta$ 的影响。
• 推导出两个关键时间尺度：$\tau_z \sim \|\mathbf{Z}^0\|_2 / \tilde{\eta}$ 表示学习启动时间，$\tau_{nl} \sim \beta / \tilde{\eta}$ 表示非线性偏离时间。
• 采用相关初始化策略，在保持神经正切核 $\hat{\Theta}$ 不变的前提下，独立控制 $\beta$ 与 $\|\mathbf{Z}^0\|_2$。
• 在Wide ResNet（CIFAR10）、ResNet-50（ImageNet）和GRUs（IMDB）上，使用SGD和JAX/Neural Tangles进行大规模实证评估。
• 系统性地消融 $\beta$ 在 $10^{-2}$ 到 $10^1$ 范围内的影响，并分析训练稳定性、收敛速度与测试准确率。
• 测量梯度幅度与时间尺度，以验证理论预测的学习动态。

实验结果

研究问题

• RQ1 逆温度 $\beta$ 如何影响使用softmax交叉熵损失训练的模型的早期学习时间尺度？
• RQ2 泛化性能在多大程度上依赖于 $\beta$，而非初始logit幅度 $\|\mathbf{Z}^0\|_2$？
• RQ3 在非单位 $\beta$ 下泛化能力的提升，是由于模型置信度的变化，还是训练过程中的动力学效应所致？
• RQ4 在多样化架构与数据集上，最大化性能的 $\beta$ 最优范围是什么？
• RQ5 学习率如何随 $\beta$ 变化以维持稳定训练？由此得出的最优 $\beta$ 是什么？

主要发现

• 泛化性能强烈依赖于 $\beta$，最优值通常远非1，表明 $\beta$ 是关键超参数。
• 性能在很大程度上与初始logit幅度 $\|\mathbf{Z}^0\|_2$ 无关，除非 $\|\mathbf{Z}^0\|_2 \gg 1$，此时性能会下降。
• 较小的 $\beta$ 导致更快的初始学习，可能带来更高的峰值性能，但也会增加训练不稳定性，尤其在无批量归一化时。
• 最优学习率与 $\eta^* \sim 1/\beta$ 成比例，且早期学习时间尺度为 $\tau_z^* \sim \|\mathbf{Z}^0\|_2 / \beta$。
• 实证结果表明，在 $[10^{-2}, 10^1]$ 范围内调节 $\beta$ 可获得一致的性能提升，且在所有架构中并无单一最优值。
• 泛化对 $\beta$ 的依赖是一种动力学现象，而非模型置信度变化所致，该结论通过消融分析与校准分析得到验证。