Training More Robust Classification Model via Discriminative Loss and Gaussian Noise Injection

TL;DR

This work tackles the robustness-accuracy trade-off in deep classifiers by introducing a dual-objective training framework: a penultimate-layer loss that enforces intra-class compactness and maximizes inter-class margins with analytically defined decision boundaries, and a class-wise alignment of noisy data to clean data clusters. It further provides a theoretical connection: training with additive Gaussian noise implicitly reduces the eigenvalues of the input-space Hessian, i.e., the loss curvature, which regularizes robustness; reduced curvature also yields more stable features under perturbations. The authors demonstrate, across CIFAR-10, SVHN, and a road-image dataset, that their method improves robustness to diverse perturbations while preserving or even improving clean accuracy, and show that curvature reduction correlates with robustness and generalizes beyond training perturbations. Overall, the approach offers a principled, easily adoptable strategy that enhances noise robustness through feature-space regularization and curvature-aware analysis, with broad practical impact for reliable deployment of vision systems in noisy environments.

Abstract

Robustness of deep neural networks to input noise remains a critical challenge, as naive noise injection often degrades accuracy on clean (uncorrupted) data. We propose a novel training framework that addresses this trade-off through two complementary objectives. First, we introduce a loss function applied at the penultimate layer that explicitly enforces intra-class compactness and increases the margin to analytically defined decision boundaries. This enhances feature discriminativeness and class separability for clean data. Second, we propose a class-wise feature alignment mechanism that brings noisy data clusters closer to their clean counterparts. Furthermore, we provide a theoretical analysis demonstrating that improving feature stability under additive Gaussian noise implicitly reduces the curvature of the softmax loss landscape in input space, as measured by Hessian eigenvalues.This thus naturally enhances robustness without explicit curvature penalties. Conversely, we also theoretically show that lower curvatures lead to more robust models. We validate the effectiveness of our method on standard benchmarks and our custom dataset. Our approach significantly reinforces model robustness to various perturbations while maintaining high accuracy on clean data, advancing the understanding and practice of noise-robust deep learning.

Figures (12)

• Figure 1: t-SNE feature visualization for test set of CIFAR10 of original (clean) and noisy data (additive Gaussian noise), produced by models trained with different methods. When training the model with the clean data, it produces features that are not well separated between the classes for noisy data (right plot of Fig. \ref{['fig:tsne_normal']}). Training model with both clean and noisy data helps to produce more discriminative features for the noisy data (Fig. \ref{['fig:tsne_noise']}) but also makes features on the clean data less discriminative. Our method helps the model to produce more discriminative features both on clean and noisy data (Fig. \ref{['fig:tsne_ours']}).
• Figure 2: Illustration of our method applied on the features of the penultimate layer. In the feature space, each color represents the data cluster of each class, where darker and lighter colors represent clean and noisy data, respectively. Our method focuses on boosting discriminativeness of features on clean data (by enforcing intra-class compactness and inter-class separability) and aligning noisy data clusters with those of clean data from the same class (in feature space).
• Figure 3: Pipeline from input $x$ to softmax output $\sigma(z)\in\Delta^{C-1}$, which is the simplex in $\mathbb{R}^C$ of probability measures. The whole softmax model is denoted by $\mathcal{N}_{\Theta}$, composed of $f_{\theta}$ and the softmax layer, where $\Theta$ is the concatenation of all the parameters of $(\theta,W,b)$.
• Figure 4: Motivating remark. Ideally, inputs from the same class produce probability vectors close to each other and near a corner of $\Delta^{C-1}$. However, for a given feature $q \in \mathcal{F}$, there exists an entire subspace mapping to the same probability vector. Consequently, using the softmax loss alone does not enforce small class dispersion.
• Figure 5: For any $\nu_1, \nu_2 > 1$, the vectors $\mathcal{T}_{\nu_1}(\Theta) - \Theta$ and $\mathcal{T}_{1/\nu_2}(\Theta) - \Theta$ point in opposite directions. Moving along either direction preserves model predictions but increases intra-class dispersion or decreases the inter-class margin, thereby violating the joint optimization enforced by our constraints. Therefore, our loss avoids such solutions $\mathcal{T}_{\nu}$, which do not improve generalization.

Theorems & Definitions (51)

• Remark 3.1
• Definition 3.1: Class dispersion
• Proposition 4.1
• Proof 4.1
• Proposition 5.1
• Proof 5.1
• Definition 5.1: Class margin
• Proposition 5.2
• Proof 5.2
• Proposition 5.3