Logit-Gap Regularization

• Logit-gap regularization is defined as a method that constrains the difference between the highest and second-highest logits to maintain a robust decision margin.
• It integrates with cross-entropy loss in adversarial training by adding a hinge-style penalty, which helps reduce overconfident predictions under perturbations.
• This technique is computationally efficient and complements methods like label smoothing and logit squeezing to improve overall adversarial defense.

Logit-gap regularization refers to a class of methods designed to enhance adversarial robustness in neural networks by explicitly constraining the difference between the largest and second-largest logits—the pre-softmax activations of the model—during training. These regularizers serve to control the decision margin, prevent overly confident (peaky) predictions, and stabilize the classifier’s response to adversarial perturbations. Logit-gap regularization is a natural extension of the broader family of logit-regularization techniques such as logit squeezing (L₂ penalty on logits), adversarial logit pairing, and label smoothing. While the top-two logit gap penalty is not the central primitive in prior work, its principles are theoretically and empirically grounded in the analysis of adversarial defense methods that operate on the logits (Summers et al., 2019).

1. Definition and Mathematical Formulation

Given a $C$-class classifier parameterized by $\theta$, denote the logit (pre-softmax activation) vector $z(x; \theta) = (z_1, ..., z_C)$. Define

• $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$ (largest logit),
• $$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$ with $c^\ast = \arg\max_c z_c(x; \theta)$ (second-largest logit).

The logit-gap regularizer is

$R_{\text{gap}}(z) = \max(0, z_{(2)} - z_{(1)}).$

Alternatively, in a margin-encouraging (hinge) form,

$$R_{\text{gap}}(z, m) = \max\left(0,\, m - [z_{(1)} - z_{(2)}]\right),$$

which with $m = 0$ reduces to the definition above. The regularizer penalizes instances where the difference between the top logit and the next largest logit is negative, i.e., where the correct class logit does not lead significantly.

The subgradient of $R_{\text{gap}}$ with respect to the logits (for input $\theta$0) is as follows:

• If $\theta$1, the gradient is zero.
• If $\theta$2, the subgradient w.r.t.\ the leading ($\theta$3) and second ($\theta$4) logit is $\theta$5 and $\theta$6 respectively; the gradient is zero for all other classes.

This construction ensures that the penalty is imposed only when the model’s predictions are “fragile” with respect to class ordering.

2. Integration with Training Objectives

Logit-gap regularization is incorporated into the overall training objective in combination with cross-entropy (CE) loss on clean or adversarial examples. For a mini-batch $\theta$7, the combined loss in an adversarial setting (e.g., using PGD-based adversary $\theta$8) is:

$\theta$9

where $z(x; \theta) = (z_1, ..., z_C)$0 trades off between classification fit and penalizing small margins. More generally, training can involve mixtures of clean and adversarial data, with the gap penalty optionally applied to each,

$z(x; \theta) = (z_1, ..., z_C)$1

with $z(x; \theta) = (z_1, ..., z_C)$2. The parameter $z(x; \theta) = (z_1, ..., z_C)$3 is typically chosen in $z(x; \theta) = (z_1, ..., z_C)$4, and is often warmed up gradually across the initial epochs to preserve early learning signals.

3. Theoretical Justification

Logit-gap regularization is motivated by both margin-based and Lipschitz-centric robustness frameworks:

Margin-Based View: Classical margin theory for classifiers seeks to maximize $z(x; \theta) = (z_1, ..., z_C)$5, the difference between the logit for the correct class and the closest competitor. Penalizing the negative margin (i.e., cases where $z(x; \theta) = (z_1, ..., z_C)$6) compels the network to avoid ambiguous decisions and supports adversarial robustness, as small input perturbations are less likely to alter the model’s output if the margin is maintained.

Lipschitz Perspective: Limiting $z(x; \theta) = (z_1, ..., z_C)$7 enforces a tight local Lipschitz constraint on logit differences, since

$z(x; \theta) = (z_1, ..., z_C)$8

and similarly for $z(x; \theta) = (z_1, ..., z_C)$9. If the two top logits cannot cross easily under small perturbations, the argmax output remains stable, increasing adversarial resistance. Logit squeezing (L₂ penalty on logits) and adversarial logit pairing are interpreted similarly (Summers et al., 2019).

4. Algorithmic Implementation and Practical Considerations

The method is computationally lightweight and amenable to integration with existing adversarial training pipelines. The only additional computation per batch is partial sorting to find the top two logits, which is $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$0 for number of classes.

Forward Pass:

1. For each input $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$1 (clean or adversarial), compute logits $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$2.
2. Identify indices of the largest and second-largest logits.
3. Evaluate $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$3 and add $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$4 to the total loss.

Backward Pass:

• If the margin ($z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$5) is non-negative, no regularizer gradient propagates.
• If negative, apply a $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$6 gradient to the top logit and $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$7 to the second logit; all others remain unaffected.

Common stabilization heuristics include warming up $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$8 linearly from $z_{(1)}(x; \theta) = \max_c z_c(x; \theta)$9 during initial epochs, capping $$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$0 to prevent gradient explosions, and using conservative learning rate schedules. Empirically, $$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$1 in $$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$2 typically balances robustness and accuracy.

5. Empirical Evidence and Performance Proxy

Summers & Dinneen (Summers et al., 2019) did not directly evaluate the top-two logit gap penalty but conducted extensive experiments with related logit-regularization methods used as empirical proxies. Results on CIFAR-10 under $$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$3-bounded perturbations and both white-box and black-box attack settings demonstrate the gains from constraining logits.

White-box (CIFAR-10, $$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$4):

Training Method	Natural	FGSM	PGD(20)
Standard PGD (Madry)	75.8%	50.5%	45.3%
ALP	74.0%	52.6%	48.5%
LSq (logit-squeeze)	77.2%	45.4%	42.0%
LS (label smoothing)	86.9%	54.7%	34.4%
LRM (combined)	68.5%	52.8%	51.0%

A direct logit-gap penalty is expected to follow the “logit-squeeze” (LSq) curve, with optimal $$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$5 near $$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$6 recovering roughly half the improvement seen from adversarial logit pairing (ALP) over standard PGD.

6. Relationship to Other Logit-Regularization Methods

Logit-gap regularization is closely related to several established logit-based defenses:

• Logit Squeezing (LSq): Applies an L₂ norm penalty to all logits. Smoother but can reduce capacity if overapplied.
• Label Smoothing (LS): Replaces one-hot label targets by a smoothed distribution, limiting extreme logit differences via soft target assignment. Strong as a black-box defense, but brittle under extended white-box attack (long PGD or SPSA).
• Adversarial Logit Pairing (ALP): Penalizes discrepancies between logits for clean and adversarial counterparts (L₂ or dot-product). Provides gains via both distribution pairing and incidental shrinking of logit norms.
• Decoupled Logit Pairing + Squeezing (LRM): Separates logit matching from explicit L₂ regularization, allowing independent tuning of distributional and squeezing effects.

Logit-gap penalty aligns with the spirit of hinge-style losses and sparse penalties central to margin-based SVMs, directly manipulating the local margin with explicit subgradients. However, the derivative is discontinuous at the gap threshold, and some practitioners may prefer the smoothness or convexity properties of L₂-based regularizers.

7. Application Guidelines and Practical Insights

Recommended procedure for leveraging logit-gap regularization in adversarial training includes:

• Begin with a strong baseline such as PGD-10 or PGD-20 adversarial training.
• Introduce a logit-gap penalty ($$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$7, $$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$8–$$z_{(2)}(x; \theta) = \max_{c \neq c^\ast} z_c(x; \theta)$$9), warming $c^\ast = \arg\max_c z_c(x; \theta)$0 over several epochs.
• Tune $c^\ast = \arg\max_c z_c(x; \theta)$1 to recover approximately 30–50% of the robustness gain from ALP over standard PGD.
• Use smooth surrogates (e.g., $c^\ast = \arg\max_c z_c(x; \theta)$2 with $c^\ast = \arg\max_c z_c(x; \theta)$3–10) for optimization stability if needed.
• Combine with other low-cost regularizers such as label smoothing ($c^\ast = \arg\max_c z_c(x; \theta)$4), mixup, and mild L₂ squeezing ($c^\ast = \arg\max_c z_c(x; \theta)$5).
• Re-evaluate performance under both long-run PGD ($c^\ast = \arg\max_c z_c(x; \theta)$6 steps) and strong gradient-free attacks (SPSA) to ensure no gradient masking.

A plausible implication is that regularizing the logit gap enhances adversarial robustness with minimal computational overhead, and when carefully combined with other logit-based regularizers, achieves near state-of-the-art defense on vision benchmarks at negligible cost over standard adversarial training (Summers et al., 2019).