Mitigating Neural Network Overconfidence with Logit Normalization

In plain words

A classifier trained on a fixed set of categories often sounds sure of itself even when the input belongs to none of them. That makes it hard to tell when a prediction should not be trusted, and a common way to flag unfamiliar inputs is to look at how confident the model is.

The standard training loss, cross-entropy, keeps pushing the model's raw output scores to grow larger as training goes on. Larger raw scores turn into more extreme confidence, and that confidence also shows up on inputs the model has never seen. The fix is small: during training, the score vector is rescaled to a fixed length before the loss is computed. The model can then only lower the loss by changing which way the score vector points, not by making it longer. Unfamiliar inputs end up with less extreme confidence, so familiar and unfamiliar inputs are easier to separate. It is a drop-in change to the loss, and the rest of the training stays the same.

TL;DR

LogitNorm replaces cross-entropy on raw logits with cross-entropy on logits divided by their L2 norm and by a temperature tau, so every training logit vector has length 1/tau §3.2 eq. 4. The authors tie overconfidence to logit norms that keep growing during training §3.1 fig. 1. Author-reported on CIFAR-10 with WRN-40-2, the average FPR95 of the softmax score falls from 49.52% to 15.65% across six OOD sets, and ID accuracy stays close to cross-entropy tab. 1 tab. 3. Experiments cover only CIFAR-10 and CIFAR-100 §4.1.

Problem

Out-of-distribution (OOD) detection labels an input as in-distribution (ID) when a score S(x) is at or above a threshold. The threshold is usually set so that about 95% of ID data passes §2.1 eq. 1. The default score is the maximum softmax probability (MSP) §1. The paper argues that deep networks can give abnormally high softmax confidence even for inputs far from the training data, which weakens MSP §1. Earlier work mostly changed the score function (ODIN, energy, GradNorm, KNN-based scores) §1 §5.

The setting is a supervised k-class problem with an i.i.d. training set. OOD inputs have labels outside the training label set §2.1. Metrics are FPR95 (false positive rate on OOD data at 95% true positive rate on ID data), AUROC and AUPR §4.1.

Method

Motivation §3.1. Write the logit vector as its length times its unit direction eq. 2. For a correctly classified training example, the cross-entropy loss can be lowered by making the length larger, because the softmax probability of the correct class then rises. Training therefore keeps increasing logit lengths, for ID and OOD inputs alike fig. 1. Two propositions in §3.1 show that scaling logits up keeps the argmax and raises the softmax value of the predicted class §3.1.

Objective §3.2. LogitNorm applies cross-entropy to the normalized logit vector f / ||f|| §3.2. The authors state the goal as a constraint, keeping the L2 norm of the logits fixed, and note that adding it through a Lagrange multiplier did not work well in their tests §3.2 §5. They therefore use the normalized loss instead. The explicit form has a temperature tau eq. 4:

\mathcal{L}_{\text{logit\_norm}}(f(x;\theta), y) = -\log \frac{e^{f_y/(\tau\|f\|)}}{\sum_{i=1}^{k} e^{f_i/(\tau\|f\|)}} eq. 4

Because the scaled vector has length 1/tau, the loss can only change through direction §3.2. Viewed another way, tau times the norm acts as an input-dependent temperature §3.2. A small positive constant (the paper gives 1e-7) is added to the denominator for numerical stability §3.2.

Training. The rest of the recipe is ordinary supervised training with the hyperparameters in the table below §4.1. The paper says the change needs no architecture or schedule changes §1 §5.

# LogitNorm training loss  [§3.2] [eq. 4]  [code: common/loss_function.py#L47-L57]
# Input:  logits (B, k); labels (B,); temperature tau > 0
# Output: scalar loss

def logit_norm_loss(logits, labels, tau):
    # Euclidean length of each logit vector  [§3.2]
    lengths = norm(logits, axis=-1, keepdims=True)     # (B, 1)
    # Unit direction, with a small epsilon for stability  [§3.2]
    directions = logits / (lengths + 1e-7)             # (B, k)
    # Every vector now has length 1/tau  [§3.2]
    scaled = directions / tau                          # (B, k)
    # Ordinary cross-entropy on the scaled vectors  [§3.2] [eq. 4]
    return mean(cross_entropy(scaled, labels))

Test time §4.2. The paper keeps the scoring function unchanged, so that LogitNorm's effect is isolated. The primary score is MSP computed on the raw network output, with no normalization.

# Test-time score: maximum softmax probability  [§4.2]
# Input:  raw logits of the trained network, shape (N, k)
# Output: score per input, higher for more in-distribution  (N,)

def msp_score(logits):
    # Ordinary softmax over the unmodified logits  [§4.2]
    probabilities = softmax(logits)                    # (N, k)
    # The largest class probability is the score  [§4.2]
    return max(probabilities, axis=-1)                 # (N,)

The paper also applies LogitNorm-trained networks with ODIN, energy and GradNorm scores §4.2 tab. 2. Those scores are computed as in their original papers, with the parameters given in the hyperparameters table.

Hyperparameters

Results

Author-reported. Setup: in-distribution CIFAR-10 or CIFAR-100, WRN-40-2 unless stated, six OOD test sets (Textures, SVHN, Places365, LSUN-Crop, LSUN-Resize, iSUN), softmax confidence score, five seeds averaged §4.1 tab. 1. Each cell is cross-entropy / LogitNorm.

Table 1: softmax score, cross-entropy vs LogitNorm tab. 1

Table 2: other OOD scores, average over the six sets tab. 2

Table 3: architectures, CIFAR-10 ID tab. 3

Table 4: expected calibration error, %, WRN-40-2, M = 15 tab. 4

Table 5: CIFAR-10 loss comparison, SVHN as OOD tab. 5

Table 6: GODIN vs LogitNorm, CIFAR-10 tab. 6

Headline numbers, all author-reported: the SVHN FPR95 falls from 50.33 to 8.03 on CIFAR-10 §4.2 tab. 1; ODIN's FPR95 falls from 40.32 to 12.95 tab. 2; on DenseNet the average FPR95 falls from 50.41 to 18.57 tab. 3. The text also gives CIFAR-100 accuracy of 75.12% for LogitNorm and 75.23% for cross-entropy on ResNet-34, which appears in no table §4.2. Per-dataset tables for the other scores and architectures are in the appendix as tab. 7, tab. 8 and tab. 9 and are not copied here.

Paper vs code

Code: the paper links https://github.com/hongxin001/logitnorm_ood §1. The repository's README says it is the official implementation [code: README.md#L4]. Checked at commit 0a60eeffb7dfc970fe68e07c5649ea1c9c8244c6.

• Temperature default. Paper: tau = 0.04 by default for CIFAR-10 §4.1. Code: --temp defaults to 1.0 train.py. README example: --temp 0.01 README.md, which does not name a dataset. Paper and README disagree, and the repository does not record which tau produced the shipped checkpoint.
• Tuning set. Paper: tau is tuned on Gaussian noise §4.1. Code: no noise set. --validation holds out 10% of the training data by index train.py validation_dataset.py. That split is in-distribution, so it cannot reproduce the paper's protocol.
• Seeds. Paper: five repetitions, averaged §4.1. Code: a single seed, default 1 train.py. No repetition loop exists.
• Augmentation. Not stated in the paper. Code: random horizontal flip and random 32 x 32 crop with 4-pixel padding, then CIFAR mean and std normalization utils.py.
• ODIN and energy parameters. Paper: ODIN T = 1000, epsilon = 0.0014 §4.2; energy T = 0.1 (CIFAR-10) and 0.01 (CIFAR-100) for LogitNorm §4.2. Code: --T and --noise default to 1 and 0 test.py, and the repository ships no per-method evaluation commands.
• Sign convention. Paper: higher score means in-distribution §2.1. Code: MSP and energy are returned negated, so lower means more in-distribution ood_tools.py. Check the convention before comparing numbers.
• Epsilon. Paper: epsilon added to the denominator §3.2. Code: added to the norm before division loss_function.py. The two are the same formula for practical purposes.
• Loss, optimizer, schedule, batch, epochs, dropout, weight decay. These match the paper. Loss: loss_function.py. Optimizer and milestones: base_framework.py. Batch, epochs, dropout, weight decay: train.py.
• Shipped snapshot. snapshots/standard/ holds a CIFAR-10 WRN LogitNorm run. Its log ends at 94.50% test accuracy at epoch 200 cifar10_wrn_logitnorm_standard_training_results.csv. That is a single run, not the paper's mean over five, and its tau is not recorded.
• Data paths. The Textures and SVHN loaders hard-code /data/ood_test utils.py. Other OOD sets are not loaded in the part of that file read for this card.

Pitfalls

• Percentages are percentage points. The abstract's "42.30%" is 50.33 minus 8.03 tab. 1. The average of 33.87 is 49.52 minus 15.65 tab. 1. These are absolute differences in FPR95, not relative reductions. Relative changes would be much larger.
• CIFAR-100 averages disagree. The LogitNorm column of the CIFAR-100 softmax average in tab. 1 is 58.49 / 85.73 / 96.42. The softmax row of tab. 2 gives 58.35 / 85.75 / 96.44. Treat the last digit as uncertain.
• Caption of tab. 6 does not match its numbers. The caption says LogitNorm uses the ODIN score for a fair comparison. The LogitNorm row (15.65 / 96.91 / 99.31) equals the MSP average in tab. 1. The text says MSP was used §5. ODIN with LogitNorm is 12.95 FPR95 in tab. 2.
• Caption of tab. 4 does not match its rows. The caption names only CIFAR-10, but the rows include CIFAR-100.
• Uncommented calibration result. Without temperature scaling, LogitNorm's ECE is 66.95% on CIFAR-10 against 3.20% for cross-entropy tab. 4. The text does not discuss this. The calibration claim rests on the "with temperature scaling" rows §4.2.
• Scale of the logits. At training, each logit vector has length 1/tau §3.2. The proof bounds each coordinate by 1/tau in magnitude app. proofs of §3.2. With tau = 0.04, logits fed to the softmax lie in about [-25, 25]. The paper does not discuss this range.
• Loss floor is essentially zero at tau = 0.04. The lower bound log(1 + (k-1) e^(-2/tau)) §3.2 gives about 1.7e-21 for k = 10 and tau = 0.04. This is my computation from the paper's bound. The paper says a large bound hurts optimisation, which is why small tau is preferred §3.2.
• Tuning cost. Choosing tau needs several training runs. The paper calls this computationally expensive app. Future Work.
• Tuning protocol. The paper tunes on Gaussian noise, which is not in the repository (see Paper vs code). Results may depend on this choice, and the paper does not report the sensitivity.
• Logit-penalty baseline. The penalty alternative can fail to converge when lambda is large, and the paper says a large lambda is needed to regularize the norm §5. Lambda is 0.05 in the appendix tab. 7, not in the main text.
• Inference uses raw logits. LogitNorm changes only the training loss. Normalizing logits at test time would change the scores, and the paper does not test that.
• No variance reported. Tables give averages over five seeds with no standard deviation, so gaps of a few tenths of a point are not interpretable.
• Scope. Experiments use CIFAR-10 and CIFAR-100 only §4.1. No larger datasets are reported.

Lineage

• Builds on: maximum softmax probability (MSP), ODIN, the energy score and GradNorm, which are the scoring functions the paper applies to LogitNorm-trained networks §1 §4.2.
• Compared against: GODIN and cosine-similarity losses, which normalize weights and features §5; logit penalty, an L2-norm regularizer §5.
• Related to temperature scaling: LogitNorm's effective temperature is input-dependent and is set during training, while ODIN's is one constant at test time §5.