Towards Better Generalization via Distributional Input Projection Network
Extracted source text · arXiv 2506.04690
\documentclass{article}
\input{myheaders}
% if you need to pass options to natbib, use, e.g.:
% \PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2025
% ready for submission
\usepackage[preprint]{neurips_2025}
% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
% \usepackage[preprint]{neurips_2025}
% to compile a camera-ready version, add the [final] option, e.g.:
% \usepackage[final]{neurips_2025}
% to avoid loading the natbib package, add option nonatbib:
% \usepackage[nonatbib]{neurips_2025}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
\usepackage{hyperref} % hyperlinks
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{amsfonts} % blackboard math symbols
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{microtype} % microtypography
\usepackage{xcolor} % colors
\usepackage{enumitem}
\newcommand\xinwei[1]{{\color{cyan}Xinwei: #1}}
\newcommand{\name}{DIPNet}
%\title{Towards Better Generalization by Learnable Smoothing on Neural Networks}
\title{Towards Better Generalization via Distributional Input Projection Network}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.
\author{%
Yifan Hao$^1$\thanks{Random order, equal contribution.\\ Emails: \texttt{yifanh12@illinois.edu, yanxinl4@illinois.edu, xinwei.shen@stat.math.ethz.ch, tongzhang@tongzhang-ml.org}} \quad \quad Yanxin Lu$^1$$^*$ \quad \quad Xinwei Shen$^2$\quad \quad Tong Zhang$^1$\\
$^1$University of Illinois Urbana-Champaign\\
$^2$Eidgenossische Technische Hochschule Zurich Departement Mathematik\\
% \texttt{hippo@cs.cranberry-lemon.edu} \\
% examples of more authors
% \And
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email} \\
% \AND
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email} \\
% \And
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email} \\
% \And
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email} \\
}
\begin{document}
\maketitle
\begin{abstract}
As overparameterized models become increasingly prevalent, training loss alone offers limited insight into generalization performance. While smoothness has been linked to improved generalization across various settings, directly enforcing smoothness in neural networks remains challenging. To address this, we introduce Distributional Input Projection Networks (DIPNet), a novel framework that projects inputs into learnable distributions at each layer. This distributional representation induces a smoother loss landscape with respect to the input, promoting better generalization.
We provide theoretical analysis showing that DIPNet reduces both local smoothness measures and the Lipschitz constant of the network, contributing to improved generalization performance. Empirically, we validate DIPNet across a wide range of architectures and tasks, including Vision Transformers (ViTs), Large Language Models (LLMs), ResNet and MLPs. Our method consistently enhances test performance under standard settings, adversarial attacks, out-of-distribution inputs, and reasoning benchmarks.
We demonstrate that the proposed input projection strategy can be seamlessly integrated into existing models, providing a general and effective approach for boosting generalization performance in modern deep learning.
%We also evaluate its Out-Of-Distribution (OOD) performance across multiple datasets, showcasing improved OOD robustness of such method. Additionally, extending our method to perturb ResNet across all blocks yields enhanced test accuracy. All of these findings highlight the effectiveness of our algorithm, in the sense that ``smoothing'' the neural network can lead to improved generalization performance.
% Focusing on this issue, we consider perturb the training data with random samples from a learnable distribution. This perturbation can lead to smaller Lipschitz norm on neural networks, corresponding to a better generalization performance. We provide a comprehensive empirical analysis of our algorithm, by showing that it leads to a meaningful decreasing on test loss in many data sets and neural network structures. We also verified its advantages on forecasting accuracy, comparing with general adversarial training methods, while both of them could achieve good adversarial robustness. Furthermore, we investigate its Out-Of-Distribution (OOD) performance on several datasets, establishing its improvement on OOD robustness. And we also extend to disturb Resnet on its all blocks, which also leads to an improvement on test accuracy. All of these results highlight the effectiveness of our algorithm, in the sense that ``smoothing'' the neural network can lead to improved generalization performance.
\end{abstract}
\section{Introduction}
Overparameterization has become a defining feature of modern deep learning. From vision transformers to large language models, today's networks often possess far more parameters than training examples. While this overparameterization enables remarkable expressivity and optimization ease, it is of great importance to identify models with strong generalization performance (i.e, the excellent performance on unseen test data).
Recent work has sought to characterize generalization through various geometric and functional properties of the learned models. Notably, \citet{johnson2023inconsistency} argue that the generalization gap can be largely attributed to two components: inconsistency and instability, with the latter being closely related to the Lipschitz continuity of the learned function. In parallel, another prominent line of research focuses on sharpness, typically quantified via the spectral norm of the Hessian. Building on this perspective, \citet{foret2020sharpness} proposed Sharpness-Aware Minimization (SAM), a technique aimed at locating flatter regions of the loss landscape. SAM has been successfully applied in diverse domains including computer vision \citep{chen2021vision}, natural language processing \citep{bahri2021sharpness}, and bi-level optimization \citep{abbas2022sharp}.
However, enforcing low sharpness, or low Lipschitz constants during training remains a significant challenge—particularly without sacrificing generalization. For instance, there is no guarantee that a model with low Lipchitz on the training distribution will remain a low Lipschitz on the test distribution. Adversarial training \citep{madry2017towards} and Random Smoothing \citep{cohen2019certified} have proven effective in reducing the Lipschitz constant and improving robustness, but often introduces a trade-off between robustness and standard generalization performance \citep{tsipras2018robustness, zhang2019theoretically, javanmard2020precise, donhauser2021interpolation, dobriban2023provable, hao2024surprising}. Likewise, while SAM can find flatter regions of the landscape, it does not always yield true minimizers, even in sufficiently flat neighborhoods \citep{yue2023sharpness}.
An increasingly promising direction involves promoting smoothness in the learned models. Some recent works \citep{johnson2023inconsistency, shen2023engression} suggest that smoother input-output mappings lead to improved robustness and generalization. Smoothness mitigates the model’s sensitivity to small input perturbations, thereby enhancing stability across data distributions. Nonetheless, directly enforcing smoothness in deep networks remains difficult due to their nonlinear and high-dimensional nature.
In this paper, we introduce Distributional Input Projection Networks (DIPNet)—a novel architectural framework that improves generalization by implicitly promoting smoothness. Rather than passing deterministic inputs through the network, DIPNet projects inputs into learnable distributions at each layer, enabling the model to reason over localized neighborhoods in the input space. This structured stochasticity acts as a form of regularization, smoothing the loss landscape and mitigating sensitivity to input variation. As shown in Figure~\ref{fig:dipnet}, the proposed framework consistently improves generalization across a wide range of datasets and model architectures, with particularly notable gains under adversarial attacks and distribution shifts.
Moreover, DIPNet is broadly applicable and can be integrated into existing architectures with minimal overhead. Its distributional nature also draws inspiration from variational inference, which motivates our efficient training procedure with a stability-promoting penalty.
Our main contributions are as follows:
\begin{enumerate}[leftmargin=*]
\item We introduce DIPNet, a novel architectural framework that enhances smoothness by projecting inputs into learnable distributions at each layer of the network.
\item We develop an efficient training method for DIPNet, inspired by variational inference, which also incorporates a stability penalty to promote robustness and regularization.
\item We provide theoretical guarantees showing that DIPNet reduces function smoothness measures and Lipschitz constants, offering insight into its generalization benefits.
\item We conduct extensive experiments across a diverse set of tasks and architectures—including ViTs, LLMs, MLPs, ResNets, adversarial robustness settings, out-of-distribution (OOD) generalization, and reasoning benchmarks—demonstrating consistent improvements in test-time performance. Such improvements are particularly significant under adversarial attacks or distribution shifts.
\end{enumerate}
%In real-world scenarios, over-parameterized neural networks often find solutions that perform optimally or near-optimally on training data, but these solutions do not always generalize well to test data.
%Recent research has extensively demonstrated a connection between generalization performance and the smoothness of the estimated function. Specifically, better generalization is often associated with more smoothness. Furthermore, \citet{johnson2023inconsistency} recently proposed that the generalization gap can be largely dominated by the "inconsistency" and "instability" terms, with the latter being related to the function's Lipschitz. Another line of work focuses on sharpness, i.e, the spectrum norm of Hessian matrix.
%Building on this perspective, \citet{foret2020sharpness} introduced Sharpness-Aware Minimization (SAM), a method designed to yield flatter solutions, which has found applications in various domains such as CV \citep{chen2021vision}, NLP \citep{bahri2021sharpness}, and bi-level learning \citep{abbas2022sharp}.
%However, all of these characteristics are challenging to enforce without compromising generalization performance. For instance, there is no guarantee that a solution smooth on the training data will remain smooth on test data. Additionally, adversarial training \citep{madry2017towards} is an effective method for reducing the Lipschitz norm and enhancing robustness to adversarial attacks. Nonetheless, there is often an unavoidable trade-off between generalization accuracy and adversarial robustness (function Lipschitz) in many scenarios.\citep{tsipras2018robustness, zhang2019theoretically, javanmard2020precise, donhauser2021interpolation, dobriban2023provable, hao2024surprising}. Furthermore, while SAM can produce a flat solution during training, there is a risk that this solution may not be a minimizer, even if the surrounding landscape is sufficiently flat \citep{yue2023sharpness}.
%In this work, we propose a straightforward yet effective method to reduce the and hence improve the function smoothness, thereby enhancing generalization performance. Specifically, we perturb the input covariates with samples drawn from a distribution during the training process, a technique we refer to as a ``smoothing'' procedure, while simultaneously penalizing the variance of these perturbations to guarantee the consistency of the prediction model. This approach leads to significant improvements in generalization across various datasets and neural network architectures. Additionally, the improvement is further amplified when the perturbation distribution is made learnable.
\begin{figure}
\centering
\includegraphics[width=0.8\linewidth]{fig_2025/dipnet.pdf}
\caption{Pipeline of Distributional Input Projection Network.}
\label{fig:dipnet}
\end{figure}
\section{Related work}
\paragraph{Influential factors on generalization performance.}
In real-world scenarios, over-parameterized neural networks often find solutions that perform optimally or near-optimally on training data, but these solutions do not always generalize well to test data.
Given the significant challenges associated with generalization performance in deep neural networks, many studies have aimed to identify the key factors influencing generalization. From the parameter perspective, a considerable body of research has focused on the relationship between function sharpness and generalization \citep{hochreiter1997flat, keskar2016large, izmailov2018averaging, jiang2019fantastic, foret2020sharpness}. For instance, \citet{foret2020sharpness} intrtoduced a method named \emph{sharpness-aware-minimization} (SAM), which could reduce the sharpness of output model and subsequently improves generalization. However, even if the solution of SAM always aligns with flat loss landscape, it may not be optimal \citep{yue2023sharpness}.
From the input covariates perspective, aligning with the conventional smoothness viewpoint, \citet{johnson2023inconsistency} recently suggested different indicators to measure generalization, i.e, ``inconsistency'' \citep{nakkiran2020distributional, jiang2021assessing, kirsch2022note} and ``instability'' \citep{bousquet2002stability, shalev2010learnability}, with the latter being related to the Lipschitz norm of the output models. While adversarial training \citep{madry2017towards, cohen2019certified, salman2019provably, lecuyer2019certified, gowal2020uncovering, wang2021convergence, zou2021provable} can control the Lipschitz norm, it often involves a trade-off between generalization performance and adversarial robustness (Lipschitz) \citep{tsipras2018robustness, zhang2019theoretically, javanmard2020precise, donhauser2021interpolation, dobriban2023provable, hao2024surprising}. In contrast, our proposed method enhances generalization performance by projecting the input into a learnable distribution, thereby avoiding such trade-off typically observed with other approaches.
\paragraph{Smoothing methods.} Prior work has explored enhancing smoothness through perturbation-based sampling in various contexts. For example, \citet{shen2023engression} proposed Engression, a method that learns distributions via pre-additive perturbations. Similarly, diffusion models \citep{ho2020denoising, song2020denoising, dhariwal2021diffusion, saharia2022palette, rombach2022high} can be interpreted as techniques for modeling sample distributions from Gaussian noise. Another related line of research focuses on improving model robustness. In particular, certified robustness methods \citep{cohen2019certified, salman2019provably, lecuyer2019certified, yang2020randomized} aim to ensure stability under perturbations by injecting noise from specific distributions during training. Gaussian noise injection has also been widely used in data augmentation \citep{moreno2018forward} to enhance generalization by artificially expanding the training set.
While our method also seeks to promote smoothness, it differs in two key ways:
(i) Rather than treating noise injection as a training-only technique, we incorporate distributional projection directly into the model architecture, applying it consistently during both training and inference. This architectural integration enables theoretical guarantees for improved generalization.
(ii) Instead of limiting perturbations to the input layer, we project the input into a learnable distribution at every layer of the network. This layerwise control provides enhanced stability and helps mitigate issues such as gradient explosion during training.
%we consider the use of perturbations not only during training but also during the testing process. This approach provides theoretical guarantees for improved generalization. Additionally, we introduce two key extensions: (i) incorporating a penalty term related to perturbation variance, and (ii) extending the procedure to include learnable perturbation distributions.
\section{Problem Formulation and Variational Inference}
In this section, we formally define the problem. Our goal is to learn a model parameterized by $\theta$ that minimizes the negative log-likelihood loss for prediction tasks:
\begin{equation*}
\mathcal{L}(\theta) = - \E_x \ln \Pb (y | x, \theta).
\end{equation*}
\paragraph{Derivation of distributional input.} Inspired by \citet{shen2023engression}, on each layer of the model, we consider to project input $x \in \mathrm{R}^p$ into a distribution $\mathcal{N}(x, \varSigma)$, where $\varSigma$ is the learnable variance, to take a better prediction. With a training set containing $n$ samples $\{ x_i, y_i\}_{i=1}^n$, we can assume there exists an unobserved variable $\eta \sim \mathcal{N}(0, \gamma I_p)$, such that
\begin{equation}\label{eq:func}
\Pb(y | x, \eta, \theta) = \Pb(y| x+\eta, \theta).
\end{equation}
Following the standard variational‑inference derivation, we obtain the bound for $\mathcal{L}(\theta)$ as:
\begin{align*}
\mathcal{L}(\theta) &= - \sum_{i=1}^n \ln \Pb (y_i | x_i, \theta) = - \sum_{i=1}^n \ln \E_\eta \Pb (y_i, \eta | x_i, \theta) \\
&= - \sum_{i=1}^n \ln \E_\eta \left( \Pb(y_i | x_i+\eta, \theta) \Pb(\eta ) \right) \\
&\le - \sum_{i=1}^n \E_{\eta \sim q(\eta)} \left[ \ln \Pb(y_i | x_i+\eta, \theta) + \ln \Pb(\eta) - \ln q(\eta) \right],
\end{align*}
for any distribution $q(\eta)$,
where the third equality is from Eq.~\eqref{eq:func}, and the last inequality is from ELBO lower bound\footnote{Here is an upper bound as we consider the negative log-likelihood function.}. For simplicity, we constrain $q(\eta)$ to be a Gaussian distribution:
\begin{equation*}
q(\eta) := \mathcal{N}(0, \varSigma), \quad \text{where} \quad \varSigma = \mathrm{diag}\{ \lambda_1, \dots, \lambda_p \},
\end{equation*}
then the lower bound should be:
\begin{align*}
&\quad - \sum_{i=1}^n \E_{\eta \sim \mathcal{N}(0, \varSigma)} \left[ \ln \Pb(y_i | x_i+\eta, \theta) - \frac{p \ln \gamma}{2} + \frac{\ln | \varSigma|}{2} - \frac{\eta^T\eta}{2 \gamma} + \frac{\eta^T \varSigma^{-1}\eta}{2} \right] \\
&= - \sum_{i=1}^n \E_{\eta \sim \mathcal{N}(0, \varSigma)} \left[ \ln \Pb(y_i | x_i+\eta, \theta) - \frac{p \ln \gamma}{2} + \frac{\sum_{j=1}^p \ln \lambda_j}{2} - \frac{\sum_{j=1}^p \lambda_j}{2 \gamma} + \frac{p}{2} \right].
\end{align*}
Based on this formulation, to achieve accurate predictions, we minimize the following loss function:
\begin{equation}\label{eq:loss_1}
\min_{\theta, \varSigma} \left\{ - \sum_{i=1}^n \E_{\eta \sim \mathcal{N}(0, \varSigma)} \ln \Pb(y_i | x_i+\eta, \theta) + \alpha \sum_{j=1}^p \lambda_j - \beta \sum_{j=1}^p \ln \lambda_j \right\},
\end{equation}
with tuning parameter $\alpha, \beta > 0$.
\begin{remark}
The term $\alpha \sum_{j=1}^p \lambda_j - \beta \sum_{j=1}^p \ln \lambda_j$ can be regarded as a special penalty term to prevent the corresponding parameters $\{ \lambda_1, \dots, \lambda_p\}$ shrinking to zero during training process.
\end{remark}
\section{\name: Distributional Input Projection Network}
Building on the framework above, we propose a new algorithm—Distributional Input Projection Network (DIPNet), which projects the input into a learnable distribution at each layer of the neural network. An overview of the DIPNet pipeline is presented in Figure~\ref{fig:dipnet}. We describe the key components and implementation details in this section.
% In this section, we introduce our proposed algorithm detailedly.
\subsection{Learning Distributional Input Layerwise to Minimize Loss Function}
%In real-world scenarios, over-parameterized neural networks often find solutions that perform optimally or near-optimally on training data, but these solutions do not always generalize well to test data. To achieve good generalization performance, it is desirable to obtain smooth output models through the training process \citep{foret2020sharpness}.\xinwei{I feel this sentence before should appear earlier in the introduction. Here we should go straight to technical stuff.}
Motivated by enhancing generalization performance, \name\ can switch the standard multi-layer neural network into a distributional input projection framework. To be specific, on each layer of the model, denoting the output of the previous layer as $v$, \name\ performs the following steps:
\begin{enumerate}
\item Project $v$ into a learnable Gaussian distribution $\mathcal{N}(v, \varSigma)$;
\item Sample a particle $u$ from $\mathcal{N}(v, \varSigma)$;
\item Use $u$ as the input to the current layer and compute its output.
\end{enumerate}
Since each layer introduces stochasticity through distributional projection, the final output becomes inherently random. To ensure stable and reliable predictions, this input-output procedure is repeated $k$ times. The final output is then obtained by distilling the results from all $k$ sampled trajectories.
%Across multiple layers, the output on the last layer should be random due to the distributional input projection. To achieve a more stable prediction, we repeat this input-output progress $k$ times, and the final output is obtained by distillation across all $k$ random outputs.
\paragraph{Adding a stability penalty.} While \name\ can enforce smoothness by distributional projection, it induces instability on model output $f(x, \eta_1, \dots, \eta_L, \theta)$ ($L$ is the number of model layers). To address this issue, we propose to penalizing the variance of the model output. To be specific, based on Eq.~\eqref{eq:loss_1}, we now formulate the loss function as:
\begin{equation}\label{eq:loss_2}
\begin{aligned}
& \quad \min_{\theta, \{\varSigma_l\}_{l=1}^L} \mathcal{L}_{\alpha, \beta,\lambda}(\theta, \varSigma_1, \dots, \varSigma_L) := \min_{\theta, \{\varSigma_l\}_{l=1}^L} \left\{ - \sum_{i=1}^n \E_{\eta_1, \dots, \eta_L} \ln \Pb(y_i | x_i, \eta_1, \dots, \eta_L, \theta) \right. \\
& \quad \quad \left. + \alpha \sum_{l=1}^L \sum_{j=1}^{p_l} \lambda_j^l - \beta \sum_{l=1}^L \sum_{j=1}^{p_l} \ln \lambda_j^l + \underbrace{ \lambda \sum_{i=1}^n \V_{\eta_1, \dots, \eta_L} \left[ f(x_i, \eta_1, \dots, \eta_L, \theta) \right]}_{\text{stability penalty}} \right\},
\end{aligned}
\end{equation}
where $\lambda$ is a regularization parameter, and for each $l \in [L]$, we have $\eta_l \sim \mathcal{N}(0, \varSigma_l)$, and $\varSigma_l = \text{diag}\{ \lambda_1^l, \dots, \lambda_{p_l}^l \}$. With a proper choice of $\lambda$, adding such penalty in training process can avoid extremely instable model output, thereby benefits generalization \citep{johnson2023inconsistency}.
\begin{remark}
If we consider a simple one-layer linear model, i.e,
\begin{equation*}
\Pb( y | x+\eta, \theta ) \sim \mathcal{N}(\theta^T(x+\eta), \sigma^2 ),
\end{equation*}
with a constant $\sigma > 0$. The stability penalty can be expressed as
\begin{equation*}
\lambda \sum_{i=1}^n \V_{\eta \sim \mathcal{N}(0, \varSigma)} [\theta^T(x_i + \eta)] = \lambda n \theta^T \varSigma \theta,
\end{equation*}
which is equivalent to a $\ell_2$-norm penalty on model parameter.
\end{remark}
\paragraph{Unbiased estimation on loss function.} Before introducing the training algorithm formally, we first formulate the true loss function we use in practice:
\begin{equation}\label{eq:loss_prac}
\begin{aligned}
& \quad - \frac{1}{m} \sum_{i=1}^n \sum_{j=1}^m \ln \Pb(y_i | x_i, \{ \eta_{l,i,j} \}_{l=1}^L , \theta) +\alpha \sum_{l=1}^L \sum_{j=1}^{p_l} \lambda_j^l - \beta \sum_{l=1}^L \sum_{j=1}^{p_l} \ln \lambda_j^l \\
&+ \frac{\lambda}{m(m-1)} \sum_{i=1}^n \sum_{1 \le j_1 < j_2 \le m } \| f(x_i, \{ \eta_{l,i,j_1} \}_{l=1}^L , \theta) - f(x_i, \{ \eta_{l,i,j_2} \}_{l=1}^L , \theta) \|_2^2,
\end{aligned}
\end{equation}
where $\{ \eta_{l,i,j} \}$ are i.i.d. sampled from $\mathcal{N}(0, \varSigma_l)$, and $m$ is the number of sample times. This leads to the \emph{unbiased estimator} for Eq.~\eqref{eq:loss_2}. We now formally introduce the Distributional Input Projection Network (\name ) algorithm, which is summarized in Algorithm~\ref{alg:tsmo}.%, which is induced from Eq.~\eqref{eq:smt}.
\begin{algorithm}[htb]
\caption{Distributional Input Projection Network (\name )}
\label{alg:tsmo}
\begin{small}
\begin{algorithmic}[2]
\STATE \textbf{Input:} Initial parameter $\{\theta^0, \varSigma_1^0, \dots, \varSigma_L^0 \}$, training samples $\{x_i, y_i \}_{i=1}^n$, forecasting input $x'$, output model $f$, layer number $L$, repeating time in training $m$, repeating time in prediction $k$, epochs $T$, regularization $\{ \alpha, \beta, \lambda\}$ and step size $\xi$.
\STATE \textcolor{orange}{$\triangleright$ \textbf{Training process:}}
\FOR{$t = 1, \dots, T$}
\STATE For each sample $\{ x_i, y_i \}$, repeat Algorithm~\ref{alg:prac} $m$ times, to receive its corresponding output $\{f_j^t(x_i)\}_{j=1}^m$.
\STATE Update parameter $\{\theta, \varSigma_1, \dots, \varSigma_L\}$ based on Eq~\eqref{eq:loss_prac} via gradient descent.
\ENDFOR
\STATE \textcolor{cyan}{$\triangleright$ \textbf{Prediction process:}}
\STATE Repeat Algorithm~\ref{alg:prac} $k$ times with input $x'$, then receive its corresponding output $\{ f_1(x'), \dots, f_k(x')\}$.
\STATE Calculate the final prediction based on:
\begin{equation*}
f(x') = \frac{1}{k} \sum_{r=1}^k f_r(x').
\end{equation*}
\STATE \textbf{Output:} Final prediction $f(x')$.
\end{algorithmic}
\end{small}
\end{algorithm}
\paragraph{Comparing with other methods.} Injecting Gaussian noise is a widely used technique in data augmentation \citep{moreno2018forward} and adversarial training \citep{cohen2019certified} to improve generalization. In contrast, our proposed method goes beyond traditional augmentation in two key ways:
(i) Rather than treating noise injection as a training-only strategy, we incorporate distributional projection directly into the model architecture, applying it consistently during both training and inference. This architectural integration enables theoretical guarantees for improved generalization.
(ii) Instead of injecting noise at the input level only, we project the input into a learnable distribution at each layer, providing layerwise control over perturbations. This design offers better stability and helps prevent issues such as gradient explosion during training.
% \subsection{Why perturbations lead to better generalization?}\label{sec:theory}
% As we observe a significant reduction in test error in subsequent sections, a natural question arises: why do our ``smoothing'' methods lead to improved generalization performance? The answer likely lies in two aspects: the relationship between generalization and both function Lipschitz and function smoothness.
\subsection{Theoretical Results}\label{sec:theory}
We provide some theoretical guarantees on why our approaches improve the Lipschitz and smoothness of the original model. Our analysis is focusing on one single neural network layer, which is denoted as $h(\cdot, \theta): \mathrm{R}^p \to \mathrm{R}^q$:
\begin{equation*}
h(\cdot, \theta) := [h_1(\cdot, \theta), \dots, h_q(\cdot, \theta)]^T.
\end{equation*}
\paragraph{Lipschitz norm.} \citet{johnson2023inconsistency} proposed that there are two terms, i.e., ``instability'' term and ``inconsistency'' term, which are strongly predictive of the generalization gap (the difference between the performance on the training data and unseen data). Following Theorem 1 in \citet{johnson2023inconsistency}, the ``instability'' term is related to the function Lipschitz.
The following theorem demonstrates that our distributional input projection methods can
reduce the function's Lipschitz norm compared to the original model, contributing to improved generalization (The proof is in Appendix~\ref{pf:lip}).
\begin{theorem}\label{thm:lip}
Denote
\begin{align*}
& b_r := \max_x \| \nabla_x h_r(x, \theta) \|_2 , \quad \mathcal{B}_r := \{ x \in \mathrm{R}^p | \| \nabla_x h_r(x, \theta) \|_2 > c \cdot b_r \}, \quad r \in [q],
\end{align*}
for some $0 < c < 1$. If the condition $\mu \left( \cup_r \mathcal{B}_r \right) < C $ holds, for any $\epsilon > 0$,
there exists a distribution $\mathcal{P}$, such that
\begin{equation*}
\max_{x \in \mathbb{R}^p} \| \mathbb{E}_{\eta \sim \mathcal{P} } \nabla_x h_r(x + \eta, \theta) \|_2 < (c + \epsilon) \cdot b_r, \quad \forall r \in [q].
\end{equation*}
\end{theorem}
\paragraph{Smoothness.} It has been widely observed that lower function smoothness is often associated with better generalization performance in neural networks. Our proposed methods can also enforce the smoothness condition using a simple technique, as demonstrated in Theorem~\ref{thm:smoothx}:
\begin{theorem}\label{thm:smoothx}
Denote
\begin{align*}
& s_r := \max_x \| \nabla^2_x h_r(x, \theta) \|_2 , \quad \mathcal{S}_r := \{ x \in \mathrm{R}^p | \| \nabla^2_x h_r(x, \theta) \|_2 > c \cdot s_r \}, \quad r \in [q],
\end{align*}
for some $0 < c < 1$. If the condition $\mu \left( \cup_r \mathcal{S}_r \right) < C $ holds, for any $\epsilon > 0$,
there exists a distribution $\mathcal{P}$, such that
\begin{equation*}
\max_{x \in \mathbb{R}^p} \| \mathbb{E}_{\eta \sim \mathcal{P} } \nabla^2_x h_r(x + \eta, \theta) \|_2 < (c + \epsilon) \cdot s_r, \quad \forall r \in [q].
\end{equation*}
\end{theorem}
The proof is in Appendix~\ref{pf:smoothx}. The results indicate that DIPNet can reduce the smoothness of models. In the sequel, we will show how this leads to better generalization performance. %, thereby increasing the likelihood of achieving better generalization performance.
\begin{algorithm}[htb]
\caption{\name: Practical Implementation}
\label{alg:prac}
\begin{small}
\begin{algorithmic}[2]
\STATE \textbf{Input:} Model parameter $\{\theta, \varSigma_1, \dots, \varSigma_L \}$, input $x$, output model $f$, layer number $L$.
\STATE Initialize the input as $v_0$.
\FOR{$l = 1, \dots, L$}
\STATE Sample a particle $u_l$ from $\mathcal{N}(v_{l-1}, \varSigma_l)$.
\STATE Take $u_l$ as the input of Layer-$l$.
\STATE Receive the output on Layer-$l$, and denote it as $v_l$.
\ENDFOR
\STATE \textbf{Output:} $v_L$.
\end{algorithmic}
\end{small}
\end{algorithm}
\section{Experiment}\label{sec:experiment}
We evaluate our proposed method against standard training and two common generalization methods—Sharpness‑Aware Minimization (SAM)~\citep{foret2020sharpness} and Randomized Smoothing (RS)~\citep{cohen2019certified}. Additional experiment explorations and ablation studies are in Appendix~\ref{sec:supp}.
\subsection{Exploration on Vision Transformers under Adversarial Attacks}
Compared to general non-perturbation methods, \name\ achieves higher accuracy under adversarial attacks in training process, indicating their robustness. Here we consider two types of training‐time attacks:
\begin{itemize}[leftmargin=*]
\item \textbf{Randomized Gaussian noise}: we sample $\eta \sim \mathcal{N}(0, \sigma^2 I)$ for some $\sigma > 0$, and attack the input via $x \to x + \eta$ to simulate natural distribution shifts during training.
\item \textbf{Fast Gradient Sign Method (FGSM) adversarial noise} \citep{goodfellow2014explaining}: we attack examples by $x \to x + \epsilon \cdot \mathrm{sgn}\left(\nabla_x \mathcal{L}(\theta, x, y)\right) $.
This single‐step attack efficiently approximates the worst‐case attacks within an $\ell_\infty$‐ball of radius $\epsilon$.
\end{itemize}
%While achieving satisfactory adversarial robustness, such adversarial training methods often lead to a loss in generalization performance.
By injecting either type of attack into the training inputs, we can evaluate adversarial robustness by showing the accuracy on clean test data.
The results show that \name\ enhances stability under both worst‑case attacks and random noise while preserving baseline performance. We first evaluate our method in Vision Transformers (ViTs)~\citep{dosovitskiy2020vit}, a popular architecture in computer vision known for its strong performance across classification, detection, and segmentation tasks.
\textbf{Setup.} We train three ViT variants—ViT‑Tiny (5.5M parameters), ViT‑Small (21.7M parameters), and ViT‑Base (85.8M parameters)—on the CIFAR‑100 dataset~\citep{krizhevsky2009learning}, an image classification dataset which contains 100 distinct classes of small‑scale images. Each model is started from a checkpoint pretrained on ImageNet‑21k and fine‑tuned on ImageNet‑1k, using a patch size of 16 and an input resolution of $224\times224$. During training, we inject either additive Gaussian noise ($\sigma=0.2$) or FGSM adversarial attacks ($\epsilon=0.2$) into the training inputs. At evaluation, we report accuracy on clean test sets. Our implementation is adapted from the publicly available repository\footnote{\url{https://github.com/jeonsworld/ViT-pytorch}}, using default hyperparameters. Additional implementation details are provided in Appendix~\ref{sec:appendix_vit}.
\textbf{Results.} Table~\ref{tab:cifar_noise} shows the performance of various ViT models under different types of attack.
First, under the clean setting (no attack), our proposed method (\name) consistently outperforms the Standard approach across all ViT variants (ViT-Tiny, ViT-Small, ViT-Base), demonstrating improved generalization capabilities. In contrast, RS, which introduces Gaussian perturbations at the input layer, significantly decreases clean accuracy, particularly evident in the ViT-Tiny model. See Appendix~\ref{sec:appendix_vit} for training and validation curves.
Second, under Gaussian noise attacks, SAM and RS methods perform similarly to the Standard method, showing marginal improvements or even slight degradations. Conversely, our \name\ method notably enhances robustness in the ViT-Tiny and ViT-Small and remains comparable to the baseline for the ViT-Base (further details on the ViT-Base Gaussian scenario are in Section~\ref{sec:ablation}).
Finally, under FGSM adversarial noise, our approach significantly outperforms both SAM and RS methods in the ViT-Tiny and ViT-Small. This highlights the effectiveness and superiority of \name\ in mitigating adversarial attacks. However, on ViT‑Base, RS achieves higher test accuracy. This likely occurs because RS’s randomized perturbation partially attenuates the FGSM attack—while \name\ drives stronger adversarial fitting (higher train accuracy), RS can also weaken the attack and better prevent overfitting. See Appendix~\ref{sec:appendix_vit} for training and validation curves.
Overall, our method consistently improves robustness across various attacks while preserving or slightly enhancing clean accuracy, indicating strong stability and generalization capabilities.
\begin{table}[htb]
\centering
\caption{Test accuracy (\%) on CIFAR-100 under different training-time attacks.}
\label{tab:cifar_noise}
\resizebox{1\linewidth}{!}{%
\begin{tabular}{c | l c c c}
\toprule
\textbf{Attack type} & \textbf{Method} & \textbf{ViT-Tiny} & \textbf{ViT-Small} & \textbf{ViT-Base} \\
\midrule
\multirow{4}{*}{None}
& Standard & 84.71 & 89.59 & 92.46 \\
& Sharpness‑Aware Minimization (SAM) & 84.46 & 89.72 & 92.68 \\
& Randomized Smoothing (RS) & 65.45 & 84.84 & 88.32 \\
& \name & \textbf{85.36} & \textbf{90.06} & \textbf{92.84} \\
\midrule
\multirow{4}{*}{Gaussian}
& Standard & 46.31 & 75.65 & 69.13 \\
& Sharpness‑Aware Minimization (SAM) & 46.04 & 74.22 & 69.28 \\
& Randomized Smoothing (RS) & 47.04 & 74.31 & 68.74 \\
& \name & \textbf{51.62} & \textbf{78.21} & \textbf{69.31} \\
\midrule
\multirow{4}{*}{FGSM}
& Standard & 20.89 & 65.64 & 70.75 \\
& Sharpness‑Aware Minimization (SAM) & 22.63 & 66.86 & 70.34 \\
& Randomized Smoothing (RS) & 23.22 & 65.55 & \textbf{77.30} \\
& \name & \textbf{27.19} & \textbf{69.17} & 74.17 \\
\bottomrule
\end{tabular}%
}
\end{table}
\begin{table}[htb]
\centering
\caption{Flexible-extract accuracy (\%) on GSM8K.}
\label{tab:flexible_extract}
\begin{tabular}{l c c c}
\toprule
\textbf{Method} & \textbf{Qwen2.5-3B} & \textbf{Llama-3.2-3B} & \textbf{Gemma‑3‑4B‑PT} \\
\midrule
Base Model & 74.00 & 29.49 & 41.55 \\
Standard Fine‑tuning & 74.45 & 32.68 & 42.00 \\
Sharpness‑Aware Minimization (SAM) & 74.22 & 31.84 & 39.65 \\
Randomized Smoothing (RS) & 75.59 & 33.06 & 41.24 \\
\name & \textbf{75.66} & \textbf{35.56} & \textbf{42.08} \\
\bottomrule
\end{tabular}
\end{table}
\subsection{Exploration on LLM Reasoning}
We extend our method to large-scale language models (LLMs), which often contain billions of parameters and now power a wide range of applications—from machine translation and summarization to dialogue systems and deep thinking. In particular, mathematical reasoning tasks like GSM8K~\citep{cobbe2021gsm8k} have recently attracted significant attention for assessing an LLM’s ability to perform precise arithmetic and logical reasoning. We hypothesize that our method will help LLMs form smoother latent representations, improving both generalization to new problems and robustness against prompt variations.
\textbf{Setup.} We use the official OpenAI GSM8K dataset, which consists of 1,319 grade‑school math word problems requiring several reasoning steps and an exact numerical answer, replacing the original ``\verb|#### <answer>|'' marker with ``\verb|The answer is <answer>|'' and removing all calculator annotations (e.g., ``\verb|<<200|$\times$\verb|5=1000>>|'') to match the LM Evaluation Harness\citep{eval-harness}’s answer‐extraction format. To test our method across diverse architectures, we select three popular open‑source models: Qwen2.5‑3B~\citep{yang2024qwen2-5}, Llama‑3.2‑3B~\citep{dubey2024llama3}, and Gemma‑3‑4B‑PT~\citep{gemma3}. Due to computational constraints, we fine‑tune each model with LoRA for a single epoch—using a learning rate of $10^{-4}$ for Qwen2.5‑3B and $2\times10^{-4}$ for the other two—then evaluate with the LM Evaluation Harness on the GSM8K CoT task. Additional implementation details are provided in Appendix~\ref{sec:appendix_llm}.
\textbf{Results.} Table~\ref{tab:flexible_extract} reports flexible‑extract accuracy (any correct numeric output). Our method consistently outperforms other baselines on all three models, yielding nontrivial gains. Results for strict‑match accuracy are given in Appendix~\ref{sec:appendix_llm}, where we observe even larger improvements in Qwen2.5‑3B and Llama‑3.2‑3B. These findings suggest that our method not only improve numeric correctness but also enhance the model’s ability to generalize to answer formats.
% Further study could be done in a separate evaluation where the training and test answer formats differ, demonstrating enhanced robustness to format shifts and generalization.
\subsection{Out‑Of‑Distribution (OOD) Evaluation}\label{sec:ood_re}
We also explore the generalization performance on OOD tasks. We conduct experiments on tabular datasets using different neural network architectures, further validating the advantages of our proposed methods in both In-Distribution (ID) and Out-Of-Distribution (OOD) generalization performance.
\textbf{Setup.} To assess robustness under distribution shift, we use Shifts Weather Prediction
dataset \citep{malinin2021shifts}, which provides a scalar regression task to forecast the air temperature.
We train three MLP architectures—2 layers with 100 neurons per layer (2+100), 4 layers with 100 neurons per layer (4+100), and 4 layers with 400 neurons per layer (4+400)—using standard stochastic gradient descent on the MSE loss. We compare against the non‑perturbed baseline (Standard). To study performance under different data regimes, we vary the fraction of available training data (100\%, 66\%, and 33\%) and evaluate all three architectures under each split.
%We focus on the regression task and adopting the Mean Squared Error (MSE) loss as metric: $\mathcal{L}(f(\theta, x), y) := \left( f(\theta, x) - y \right)^2$.
\textbf{Results.} Figure~\ref{fig:ood} presents both ID and OOD test MSEs. The results shows that comparing with standard non-perturbation method, DIPNet consistently improve ID performance, and enhance OOD generalization significantly. More detailed results are shown in Appendix~\ref{sec:appendix_ood}.
\begin{figure}
\centering
\includegraphics[width=0.8\linewidth]{fig_2025/ood.pdf}
\caption{ID and OOD MSE on Weather Dataset.}
\label{fig:ood}
\vspace{5pt}
\raggedright
\footnotesize
\emph{Notes.} Here we train on the ID dataset, validate on both ID and OOD datasets, and test the model performance on both ID and OOD datasets (with the test OOD environment differing from the validation OOD environment).
\vspace{-4mm}
\end{figure}
\subsection{Noise‑Level Ablation: Gaussian Corruption Effects Across Model Sizes}\label{sec:ablation}
To quantify the impact of our \name\ module under varying gaussian noise levels, we evaluate both ViT‐Base and ViT‐Tiny corrupted by additive Gaussian noise with standard deviation $\sigma \in \{0.2, 0.5, 1.0\}$. Table \ref{tab:gaussian_vitb_vitt} reports test accuracy for the standard (baseline) training and for our method.
For larger model, \name\ yields consistent gains over the baseline at all noise levels. The gap widens as noise increases, demonstrating that \name\ effectively enhances robustness to heavy corruption.
For smaller model, while \name\ brings large relative gains at low to moderate noise, the absolute accuracy remains low. This indicates that ViT‑Tiny's limited capacity leads to underfitting when faced with strong noise, and suggests that further architectural capacity or noise‑specific regularization may be required.
These ablation results confirm that \name\ consistently improves noise robustness, but also highlight the interaction between module efficacy and backbone capacity under severe corruption.
\begin{table}[htb]
\centering
\caption{Evaluation under Gaussian Noise Levels ($\sigma \in \{0.2, 0.5, 1.0\}$).}
\label{tab:gaussian_vitb_vitt}
\begin{tabular}{l l c c c}
\toprule
\textbf{Model} & \textbf{Method} & \textbf{0.2} & \textbf{0.5} & \textbf{1.0} \\
\midrule
\multirow{2}{*}{ViT-Base}
& Standard & 69.13 & 54.81 & 35.02 \\
& \name & \textbf{69.31} & \textbf{54.95} & \textbf{36.08}\\
\midrule
\multirow{2}{*}{ViT-Tiny}
& Standard & 46.31 & 26.99 & 15.93 \\
& \name & \textbf{51.62} & \textbf{32.75} & \textbf{17.63} \\
\bottomrule
\end{tabular}
\end{table}
% \begin{table}[htb]
% \centering
% \caption{\(\lambda\) ablation on ViT-T with FGSM adversarial noise ($\lambda \in \{0, 10^{-4}, 10^{-3}, 10^{-2}, 10^{-1}\}$).}
% \label{tab:lambda_ablation}
% \begin{tabular}{c c c c c c}
% \toprule
% Standard & 0 & \(10^{-4}\) & \(10^{-3}\) & \(10^{-2}\) & \(10^{-1}\) \\
% \midrule
% 20.89 & \textbf{27.19} & 26.31 & 25.62 & 18.24 & 8.53 \\
% \bottomrule
% \end{tabular}
% \end{table}
\section{Conclusion}\label{sec:conclu}
In this work, we proposed Distributional Input Projection Networks (DIPNet), a novel architectural framework designed to enhance generalization by projecting inputs into learnable distributions at each layer. This approach implicitly enforces smoothness and reduces the Lipschitz constant of the network, supported by both theoretical analysis and extensive empirical validation. Across a wide range of models and tasks, DIPNet consistently improves test performance in standard, adversarial, and out-of-distribution settings. Our results demonstrate that DIPNet offers a broadly applicable and effective strategy for improving generalization in overparameterized deep networks. We consider extending this framework to reasoning tasks and reinforcement learning as a promising direction for future work.
\section*{Acknowledgements}
We would like to thank Rui Pan for his helpful comments and suggestions.
\bibliography{ref.bib}
\bibliographystyle{apalike}
\clearpage
\newpage
\appendix
\section{Theoretical Analysis}\label{pf}
\subsection{Proof of Theorem~\ref{thm:lip}}\label{pf:lip}
We can consider $\mathcal{P}$ as a uniform distribution on $\mathrm{R}^p$, which is supported on a set with measurement $\zeta C$, then for any $r \in [q]$, we can obtain that
\begin{align*}
\| \mathbb{E}_\eta \nabla_x h_r(x + \eta, \theta) \|_2^2 &\le \mathbb{E}_\eta \| \nabla_x h_r(x + \eta, \theta) \|_2^2 \\
&= \mathbb{E}_\eta \left[\| \nabla_x h_r(x + \eta, \theta) \|_2^2 | \bm{1}(x + \eta \in \cup_r \mathcal{B}_r) \right] \mathbb{P}( \cup_r \mathcal{B}_r ) \\
& \quad + \mathbb{E}_\eta \left[\| \nabla_x h_r(x + \eta, \theta) \|_2^2 | \bm{1}(x + \eta \not\in \cup_r \mathcal{B}_r) \right] \mathbb{P}( (\cup_r \mathcal{B}_r)^c ) \\
&\le \frac{1}{\zeta C} \left( C \cdot b_r + \zeta C \cdot c b_r \right) = c b_r + \frac{1}{\zeta} b_r.
\end{align*}
Choosing $\zeta > \epsilon^{-1}$, we can finish the proof.
\subsection{Proof of Theorem~\ref{thm:smoothx}}\label{pf:smoothx}
The proof is similar to the proof for Theorem~\ref{thm:lip}. We can consider $\mathcal{P}$ as a uniform distribution on $\mathrm{R}^p$, which is supported on a set with measurement $\zeta C$, then for any $r \in [q]$, we can obtain that
\begin{align*}
\| \mathbb{E}_\eta \nabla^2_x h_r(x + \eta, \theta) \|_2^2 &\le \mathbb{E}_\eta \| \nabla^2_x h_r(x + \eta, \theta) \|_2^2 \\
&= \mathbb{E}_\eta \left[\| \nabla^2_x h_r(x + \eta, \theta) \|_2^2 | \bm{1}(x + \eta \in \cup_r \mathcal{S}_r) \right] \mathbb{P}( \cup_r \mathcal{S}_r ) \\
& \quad + \mathbb{E}_\eta \left[\| \nabla^2_x h_r(x + \eta, \theta) \|_2^2 | \bm{1}(x + \eta \not\in \cup_r \mathcal{S}_r) \right] \mathbb{P}( (\cup_r \mathcal{S}_r)^c ) \\
&\le \frac{1}{\zeta C} \left( C \cdot s_r + \zeta C \cdot c s_r \right) = c s_r + \frac{1}{\zeta} s_r.
\end{align*}
Choosing $\zeta > \epsilon^{-1}$, we can finish the proof.
\section{Dataset Description}
\label{sec:dataset_description}
\subsection{Tabular Dataset for MLP}
\begin{itemize}[leftmargin=*]
\item \textbf{Air Quality.} \citep{misc_air_quality_360} The Air Quality dataset comprises 15 features derived from hourly averaged responses of an array of 5 metal oxide chemical sensors embedded in an Air Quality Chemical Multisensor Device in an Italian city, collected from March 2004 to February 2005. In this study, we use 9 of these features, with ``PT08.S3(NOx)'' as the response variable and the remaining 8 features as input variables.
\item \textbf{ETD.} \citep{haoyietal-informer-2021} The ETD dataset is related to the electronic data distribution hourly data recordings. It contains 8 features, including the date of the point, 6 different types of external power load features and the predictive value ``oil temperature''.
\item \textbf{Sharing Bike.} \citep{fanaee2014event} This dataset aims to understand the factors influencing the demand for shared bikes in the American market, with hourly data recorded from January 2011 to December 2012. For this analysis, we consider 4 features, i.e., weather, temperature, humidity, and feeling temperature, to predict the total count of rental bikes.
\item \textbf{NASDAQ100.} \citep{qin2017dual} It contains stock prices of 81 corporations under NASDAQ 100 and the index value of NASDAQ 100. The frequency of the data collection is one-minute. Here we consider the stock prices of different corporations as input variables, and the index value of NASDAQ 100 as predictive variable.
\end{itemize}
\begin{minipage}{0.52\linewidth}
%\begin{figure}
\captionsetup{width=\linewidth}
\captionof{table}{Overview of tabular dataset.}
\label{tab:dt}
\centering
\resizebox{\linewidth}{!}{
\begin{tabular}{c|cccc}
\toprule
Dataset & Air Quality & ETD & Sharing Bike & NASDAQ100 \\
\midrule
input & vector & vector & vector & vector \\
task & reg. & reg. & reg. & reg. \\
$\#$ samples & 8991 & 17420 & 17379 & 40560 \\
\bottomrule
\end{tabular}
}
% \vspace{-2mm}
%\end{figure}
\end{minipage}
\begin{minipage}{0.45\linewidth}
% \begin{table}\small
\captionsetup{width=\linewidth}
\captionof{table}{Overview of OOD tabular dataset.}
\label{tab:ood_dt}
\centering
\resizebox{\linewidth}{!}{
\begin{tabular}{c|ccc}
\toprule
Dataset & $VPower_s$ & $VPower_r$ & Weather \\
\midrule
input & vector & vector & vector \\
task & reg. & reg. & reg. \\
$\#$ samples & 546543 & 554642 & 436733 \\
\bottomrule
\end{tabular}
}
% \vspace{-2mm}
% \end{table}
\end{minipage}
\subsection{OOD Dataset}
To assess robustness under distribution shift, we use the Shifts vessel power estimation dataset \citep{malinin2022shifts} and the Shifts Weather Prediction
dataset \citep{malinin2021shifts} in OOD situations (see details in Table~\ref{tab:ood_dt} and Figure~\ref{fig:ood_des}).
%More details about the OOD datasets could be found in Appendix~\ref{sec:ooddt_info}.
\begin{itemize}[leftmargin=*]
\item \textbf{Vessel Power Estimation}\footnote{Notice that this is a developing dataset and it only contains the training data and valid data recently. So here we use the ID valid dataset to take validation, and the OOD valid dataset is used for testing.} The Shifts Vessel Power Prediction dataset contains 11 different features and one scalar target (energy utilization of cargo ships). And it contains two types of data: one is synthetic dataset ($VPower_s$), the other is real dataset ($VPower_r$).
\item \textbf{Weather Prediction}\footnote{More details can be found in \url{https://github.com/Shifts-Project/shifts}.} The Shifts Weather Prediction dataset \citep{malinin2021shifts} provides a scalar regression task to forecast the air temperature. The dataset contains 127 features and a single regression target, spanning an entire year, i.e, from September 1st, 2018, to September 1st, 2019.
\end{itemize}
\begin{figure}[h]
\centering
\begin{subfigure}[b]{0.49\linewidth}
\centering
\includegraphics[width=\linewidth]{pic/vessel_data.jpg}
% \caption{}
%\label{fig:ood_dt_vessel}
\end{subfigure}
\hfill
\begin{subfigure}[b]{0.49\linewidth}
\centering
\includegraphics[width=\linewidth]{pic/weather_data.PNG}
% \caption{}
% \label{fig:ood_dt_weather}
\end{subfigure}
\vspace{5pt}
\raggedright
\footnotesize
\emph{Notes.} The figure on the left side is about the Shifts vessel power estimation dataset, and the figure on the right right is about the Shifts Weather Prediction dataset. ``train'', ``dev'' and ``eval'' refer to training data, validation data and test data respectively. As public access to canonical ``test'' of vessel power dataset is restricted, here we valid on ``dev-in'' and use ``dev-out'' to test the model performance.
\caption{OOD dataset descriptions}
\label{fig:ood_des}
\end{figure}
\section{Supplement Experiments}\label{sec:supp}
\subsection{Explorations on MLP}\label{sec:mlp_td}
Here we compare the generalization performance of our method and a standard non-perturbation output model by evaluating them on various datasets and different neural network architectures.
To measure the generalization performance of different methods, we take explorations on four tabular datasets and detailed description could be found in Appendix~\ref{sec:dataset_description}.
% \paragraph{Air Quality} \citep{misc_air_quality_360} The Air Quality dataset comprises 15 features derived from hourly averaged responses of an array of 5 metal oxide chemical sensors embedded in an Air Quality Chemical Multisensor Device in an Italian city, collected from March 2004 to February 2005. In this study, we use 9 of these features, with ``PT08.S3(NOx)'' as the response variable and the remaining 8 features as input variables.
% \paragraph{ETD} \citep{haoyietal-informer-2021} The ETD dataset is related to the electronic data distribution hourly data recordings. It contains 8 features, including the date of the point, 6 different types of external power load features and the predictive value ``oil temperature''.
% \paragraph{Sharing Bike} \citep{fanaee2014event} This dataset aims to understand the factors influencing the demand for shared bikes in the American market, with hourly data recorded from January 2011 to December 2012. For this analysis, we consider 4 features, i.e., weather, temperature, humidity, and feeling temperature, to predict the total count of rental bikes.
% \paragraph{NASDAQ100} \citep{qin2017dual} It contains stock prices of 81 corporations under NASDAQ 100 and the index value of NASDAQ 100. The frequency of the data collection is one-minute. Here we consider the stock prices of different corporations as input variables, and the index value of NASDAQ 100 as predictive variable.
Specifically, we test different settings, including various proportions of training and test data, the number of hidden layers, and the number of neurons in each layer. Using gradient descent on the $\ell_2$ loss, we summarize the results in Table~\ref{tab:compare_sup}, indicating that both the standard errors and the adversarial errors on DIPNet are consistently smaller than those for standard networks. This further verifies the benefits of our proposed method.
%and Figure~\ref{fig:compare}.
\begin{table}[htb]
\caption{Standard and Adversarial MSE on MLP.}
\label{tab:compare_sup}
\centering
\resizebox{1.0\linewidth}{!}{
\begin{tabular}{c|c|c|cc|cc|cc}
\toprule
\textbf{Dataset} & \textbf{Test proposition} & \textbf{Method} & \textbf{St(2+100)} & \textbf{Adv(2+100)} & \textbf{St(4+100)} & \textbf{Adv(4+100)} & \textbf{St(4+400)} & \textbf{Adv(4+400)} \\
\midrule
\multirow{6}{*}{Air Quality} & \multirow{2}{*}{0.3} & Standard & 0.0858 & 24.0112 & 0.0677 & 37.1026 & 0.0628 & 20.5767 \\
& & \name & \textbf{0.0822} & \textbf{11.14} & \textbf{0.0576} & \textbf{20.9601} & \textbf{0.0526} & \textbf{20.5601} \\
\cline{2-9}
& \multirow{2}{*}{0.5} & Standard & 0.1033 & 13.8045 & 0.0873 & 49.1966 & 0.0823 & 17.7776 \\
& & \name & \textbf{0.0964} & \textbf{9.6509} & \textbf{0.0794} & \textbf{5.8675} & \textbf{0.0748} & \textbf{12.9039} \\
\cline{2-9}
& \multirow{2}{*}{0.7} & Standard & 0.0984 & 16.6161 & 0.0938 & 24.4145 & 0.0910 & 14.5242 \\
& & \name & \textbf{0.0901} & \textbf{4.8616} & \textbf{0.0844} & \textbf{6.6429} & \textbf{0.0853} & \textbf{10.1404} \\
\midrule
\multirow{6}{*}{ETD} & \multirow{2}{*}{0.3} & Standard & 0.1845 & 4.7174 & 0.1693 & 5.1239 & 0.1635 & 4.5290 \\
& & \name & \textbf{0.1787} & \textbf{3.6371} & \textbf{0.1604} & \textbf{3.7268} & \textbf{0.1578} & \textbf{2.9117} \\
\cline{2-9}
& \multirow{2}{*}{0.5} & Standard & 0.1905 & 6.1011 & 0.1724 & 4.5385 & 0.1668 & 4.5335 \\
& & \name & \textbf{0.1831} & \textbf{3.5267} & \textbf{0.1648} & \textbf{3.3667} & 0.1632 & \textbf{3.1108} \\
\cline{2-9}
& \multirow{2}{*}{0.7} & Standard & 0.2023 & 5.0696 & 0.1913 & 6.2199 & 0.1867 & 5.1136 \\
& & \name & \textbf{0.1947} & \textbf{4.3076} & \textbf{0.1776} & \textbf{2.9721} & \textbf{0.1757} & \textbf{2.9267} \\
\midrule
\multirow{6}{*}{Sharing Bike} & \multirow{2}{*}{0.3} & Standard & 0.7221 & 2.5177 & 0.7208 & 2.0264 & 0.7216 & 2.1557 \\
& & \name & \textbf{0.7188} & \textbf{2.0034} & \textbf{0.7186} & \textbf{1.8647} & \textbf{0.7182} & \textbf{1.8732} \\
\cline{2-9}
& \multirow{2}{*}{0.5} & Standard & 0.7108 & 2.3387 & 0.7107 & 2.3030 & 0.7118 & 1.8314 \\
& & \name & \textbf{0.7084} & \textbf{1.7799} & \textbf{0.7082} & \textbf{1.8780} & \textbf{0.7089} & \textbf{1.7272} \\
\cline{2-9}
& \multirow{2}{*}{0.7} & Standard & 0.7075 & 1.7905 & 0.7057 & 1.8397 & 0.7068 & 1.7931 \\
& & \name & \textbf{0.7050} & \textbf{1.7353} & \textbf{0.7043} & \textbf{1.7941} & \textbf{0.7056} & \textbf{1.7785} \\
\midrule
\multirow{6}{*}{NASDAQ100} & \multirow{2}{*}{0.3} & Standard & 3.576e-5 & 0.6066 & 3.538e-5 & \textbf{0.5447} & 2.779e-5 & 0.5303 \\
& & \name & \textbf{2.757e-5}& \textbf{0.5824} & \textbf{2.511e-5}& 0.5628 & \textbf{2.158e-5}& \textbf{0.5195} \\
\cline{2-9}
& \multirow{2}{*}{0.5} & Standard & 3.691e-5 & 0.5901 & 3.671e-5 & \textbf{0.5364} & 2.916e-5 & 0.5389 \\
& & \name & \textbf{2.890e-5}& \textbf{0.5866} & \textbf{2.642e-5}& 0.5805 & \textbf{2.266e-5}& \textbf{0.5244} \\
\cline{2-9}
& \multirow{2}{*}{0.7} & Standard & 6.146e-5 & 0.5813 & 8.085e-5 & 0.5506 & 3.612e-5 & 0.5455 \\
& & \name & \textbf{3.122e-5}& \textbf{0.5812} & \textbf{2.886e-5}& \textbf{0.5428} & \textbf{2.532e-5}& \textbf{0.5336} \\
\bottomrule
\end{tabular}
}
\vspace{5pt}
\raggedright
\footnotesize
\emph{Notes.} ``St'' ,``Adv'' refer to standard loss and adversarial loss respectively. ``2+100'', ``4+100'', ``4+400'' refer to the network architecture (e.g, ``2+100'' means the two layer network which each layer contains 100 neurons).
\vspace{-4mm}
\end{table}
%Besides, comparing with adversarial training, our proposed method could achieve good adversarial robustness while also maintaining better generalization performance simultaneously. The results are shown in Figure~\ref{fig:adv_air}, \ref{fig:adv_etd}, \ref{fig:adv_bike} and \ref{fig:adv_nas}.
% \begin{figure}[htb]
% \centering
% \begin{subfigure}[b]{0.9\linewidth}
% \centering
% \includegraphics[width=\linewidth]{pic/com_adv_air.png}
% % \caption{}
% %\label{fig:st_smo}
% \end{subfigure}
% \caption{Comparing \name\ with adversarial training on Air Quality dataset.}
% \label{fig:adv_air}
% \end{figure}
% \begin{figure}[htb]
% \centering
% \begin{subfigure}[b]{0.9\linewidth}
% \centering
% \includegraphics[width=\linewidth]{pic/com_adv_etd.png}
% % \caption{}
% %\label{fig:st_smo}
% \end{subfigure}
% \caption{Comparing \name\ with adversarial training on ETD dataset.}
% \label{fig:adv_etd}
% \end{figure}
% \begin{figure}[htb]
% \centering
% \begin{subfigure}[b]{0.9\linewidth}
% \centering
% \includegraphics[width=\linewidth]{pic/com_adv_bike.png}
% % \caption{}
% % \label{fig:st_smo}
% \end{subfigure}
% \caption{Comparing \name\ with adversarial training on Sharing Bike dataset.}
% \label{fig:adv_bike}
% \end{figure}
% \begin{figure}[htb]
% \centering
% \begin{subfigure}[b]{0.9\linewidth}
% \centering
% \includegraphics[width=\linewidth]{pic/com-adv_nas.png}
% % \caption{}
% % \label{fig:st_smo}
% \end{subfigure}
% \caption{Comparing \name\ with adversarial training on NASDAQ100 dataset.}
% \label{fig:adv_nas}
% \end{figure}
\subsection{Details for Out‑Of‑Distribution (OOD) Evaluation}\label{sec:appendix_ood}
Here we explore the generalization performance on Out-Of-Distribution (OOD) tasks. We conduct experiments on three tabular datasets using different neural network architectures, further validating the advantages of our proposed methods in OOD generalization performance.
\textbf{Setup.}\quad To assess robustness under distribution shift, we use the Shifts vessel power estimation dataset \citep{malinin2022shifts} and the Shifts Weather Prediction
dataset \citep{malinin2021shifts} in OOD situations.
For each dataset, we train three MLP architectures—2 layers with 100 neurons per layer (2+100), 4 layers with 100 neurons per layer (4+100), and 4 layers with 400 neurons per layer (4+400)—using standard stochastic gradient descent on the MSE loss. We compare against the non‑perturbed baseline (Standard). We vary the fraction of training data (100\%, 66\%, 33\%) and evaluate the same three MLP architectures under each split.
%We focus on the regression task and adopting the Mean Squared Error (MSE) loss as metric: $\mathcal{L}(f(\theta, x), y) := \left( f(\theta, x) - y \right)^2$.
\textbf{Results.}\quad Tables~\ref{tab:ood} and~\ref{tab:ood_weather} present OOD test MSEs. The OOD test error results show that \name\ consistently achieve better OOD generalization performance compared to the non-perturbation method, indicating that input perturbations enhance OOD generalization.
\begin{table}[htb]
\caption{OOD MSE on Vessel Power dataset.}
\label{tab:ood}
\centering
\resizebox{1.0\linewidth}{!}{
\begin{tabular}{c|c|ccc|ccc|ccc}
\toprule
\multicolumn{2}{c}{\textbf{\% Train}} & \multicolumn{3}{c}{\textbf{100\%}} & \multicolumn{3}{c}{\textbf{66\%}} & \multicolumn{3}{c}{\textbf{33\%}} \\
\midrule
\textbf{Dataset}& \textbf{Method} & \textbf{2+100} & \textbf{4+100} & \textbf{4+400} & \textbf{2+100} & \textbf{4+100} & \textbf{4+400} & \textbf{2+100} & \textbf{4+100} & \textbf{4+400} \\
\midrule
\multirow{3}{*}{$Vpower_s$ } & Standard & 0.0212 & 0.0231 & 0.0288 & 0.0401 & 0.0533 & 0.0601 & 0.0539 & 0.0498 & 0.0420 \\
~ & \name & \textbf{0.0195} & \textbf{0.0206} & \textbf{0.0212} & \textbf{0.0232} & \textbf{0.0305} & \textbf{0.0298} & \textbf{0.0298} & \textbf{0.0318} & \textbf{0.0324} \\
\midrule
\multirow{3}{*}{$Vpower_r$ } & Standard & 0.0400 & 0.0560 & 0.0895 & 0.0439 & 0.0401 & 0.0730 & 0.0944 & 0.0947 & 0.0859 \\
~ & \name & \textbf{0.0363} & \textbf{0.0439} & \textbf{0.0555} & \textbf{0.0345} & \textbf{0.0383} & \textbf{0.0412} & \textbf{0.0581} & \textbf{0.0534} & \textbf{0.0593} \\
\bottomrule
\end{tabular}
}
\vspace{5pt}
\raggedright
\footnotesize
\emph{Notes.} ``2+100'', ``4+100'', ``4+400'' refer to the network architecture (e.g, ``2+100'' means the two layer network which each layer contains 100 neurons).
\vspace{-4mm}
\end{table}
\begin{table}[htb]
\caption{ID and OOD MSE on Weather dataset.}
\label{tab:ood_weather}
\centering
\resizebox{1.0\linewidth}{!}{
\begin{tabular}{c|c|ccc|ccc}
\toprule
\textbf{\% Train} & \textbf{Method} & \textbf{ID(2+100)} & \textbf{ID(4+100)} & \textbf{ID(4+400)} & \textbf{OOD(2+100)} & \textbf{OOD(4+100)} & \textbf{OOD(4+400)} \\
\midrule
\multirow{3}{*}{100 \%} & Standard & 0.0360 & 0.0330 & 0.0291 & 0.0640 & 0.0647 & 0.0576 \\
~& \name & \textbf{0.0344} & \textbf{0.0322} & \textbf{0.0280} & \textbf{0.0562} & \textbf{0.0528} & \textbf{0.0547} \\
\midrule
\multirow{3}{*}{66 \%} & Standard & 0.0364 & 0.0329 & 0.0315 & 0.0898 & 0.0723 & 0.0822 \\
~& \name & \textbf{0.0354} & \textbf{0.0324} & \textbf{0.0295} & \textbf{0.0558} & \textbf{0.0552} & \textbf{0.0569} \\
\midrule
\multirow{3}{*}{33 \%} & Standard & 0.0364 & 0.0341 & 0.0325 & 0.0741 & 0.0599 & 0.0662 \\
~& \name & \textbf{0.0361} & \textbf{0.0337} & \textbf{0.0320} & \textbf{0.0519} & \textbf{0.0549} & \textbf{0.0627} \\
\bottomrule
\end{tabular}
}
\vspace{5pt}
\raggedright
\footnotesize
\emph{Notes.} Here we train on the ID dataset, validate on both ID and OOD datasets, and test the model performance on both ID and OOD datasets (with the test OOD environment differing from the validation OOD environment). ``2+100'', ``4+100'', ``4+400'' refer to the network architecture (e.g, ``2+100'' means the two layer network which each layer contains 100 neurons). ``ID'' means test data on ID environment , and ``OOD'' means test data on OOD environment.
\vspace{-4mm}
\end{table}
%\section{Out-Of-Distribution dataset descriptions}\label{sec:ooddt_info}
\subsection{Explorations on ResNet}\label{ex:res}
Furthermore, we extend our methods to ResNet architecture.
To examine its performance compared to the standard non-perturbation ResNet, we conducted experiments on CIFAR10 and CIFAR100 using different ResNet structures. We then compared the classification accuracy (in percentage) on the test data for each method. The results, summarized in Table~\ref{tab:resnet}, demonstrate the benefits of DIPNet in improving generalization performance. These findings highlighting the extended potential of our proposed methods.
\begin{table}[htb]
\caption{Result classification accuracy ($\%$) on ResNet.}
\label{tab:resnet}
\centering
\resizebox{1.0\linewidth}{!}{
\footnotesize{
\begin{tabular}{c|cccccc}
\toprule
\textbf{ Method} & \textbf{res18-cifar10} & \textbf{res34-cifar10} & \textbf{res50-cifar10} & \textbf{wideres-cifar10} & \textbf{wideres-cifar100} \\
\midrule
Standard & 90.96 & 93.52 & 93.37 & 96.14 & 80.77 \\
DIPNet & \textbf{91.72} & \textbf{93.56} & \textbf{93.89} & \textbf{96.53} & \textbf{81.73} \\
\bottomrule
\end{tabular}
}
}
\vspace{-4mm}
\end{table}
\subsection{Details for ViT Experiments}\label{sec:appendix_vit}
All Vision Transformer experiments are conducted on a single NVIDIA A100‐80G GPU. Checkpoints are loaded via the \texttt{timm} library. We enable Apex O2 mixed-precision training (FP16) with all other hyperparameters at their defaults.
For all baselines and our method, we only sample once at evaluation.
For the two baselines—SAM and RS—we first perform a hyperparameter sweep on the ViT‐Tiny model, then scale the optimal settings to larger backbones. Specifically, SAM’s perturbation radius \(\rho\) is searched over \(\{0.01,\,0.03,\,0.05\}\), yielding \(\rho = 0.05\); RS’s noise level \(\sigma\) is searched over \(\{0.001,\,0.005,\,0.01,\,0.05,\,0.1\}\), yielding \(\sigma = 0.01\).
Representative training and validation curves can be found in Figure~\ref{fig:vitt_curves}, \ref{fig:vitb_fgsm_curves}, \ref{fig:vitt_g_curves} and \ref{fig:vitt_fgsm_curves}.
\subsection{Details for LLM Experiments}\label{sec:appendix_llm}
All language‐model experiments are conducted on a single NVIDIA RTX 4090 GPU. We apply LoRA with rank 8, alpha 16, dropout 0.1, targeting the modules \texttt{["q\_proj","v\_proj"]}.
Training arguments include a batch size of 64, 1 epoch, weight decay of 0.01, warmup ratio of 0.03, and a fixed seed of 42. We search peak learning rates in \(\{7\times10^{-5},\,1\times10^{-4},\,2\times10^{-4},\,3\times10^{-4},\,5\times10^{-4}\}\), which yields \(1\times10^{-4}\) for Qwen2.5‐1.5B and \(2\times10^{-4}\) for both Llama‐3.2‐3B and Gemma‐3‐4B‐PT.
Evaluation is performed on the GSM8K CoT task with 8 shots, a batch size of 52, and seed 42.
For the two baselines—SAM and RS—we conduct hyperparameter sweeps and report the best-performing accuracy. In this case, SAM’s \(\rho\) is searched over \(\{0.01,\,0.02,\,0.05\}\) and RS’s \(\sigma\) over \(\{0.01,\,0.02,\,0.05\}\).
\begin{table}[htb]
\centering
\caption{Strict-match accuracy (\%) on GSM8K.}
\label{tab:strict_math}
\begin{tabular}{l c c c}
\toprule
\textbf{Method} & \textbf{Qwen2.5-3B} & \textbf{Llama-3.2-3B} & \textbf{Gemma‑3‑4B‑PT} \\
\midrule
Base Model & 58.83 & 27.45 & 39.95 \\
Standard Fine‑tuning & 68.61 & 31.92 & \textbf{41.55} \\
Sharpness‑Aware Minimization (SAM) & 69.37 & 30.86 & 39.27 \\
Randomized Smoothing (RS) & 69.45 & 32.22 & 40.79 \\
\name & \textbf{71.57} & \textbf{35.41} & \textbf{41.55} \\
\bottomrule
\end{tabular}
\end{table}
\subsection{Additional Ablation Study}
\begin{table}[htb]
\centering
\caption{\textbf{Ablation studies on ViT under Gaussian attack.} This table reports test accuracy (\%) on the CIFAR-100 dataset. ``Standard'' is the baseline ViT without any distributional input projection; ``Full-Layer'' (in a) and ``Learnable'' (in b) are the same learnable full distributional projection-at-every-layer setting.}
\label{tab:ablation}
\begin{subtable}[t]{\linewidth}
\centering
\caption{Depth Ablation (Layerwise DIPNet).}
\label{tab:layerwise}
\resizebox{0.52\linewidth}{!}{%
\begin{tabular}{l c c c}
\toprule
Method & ViT-Tiny & ViT-Small & ViT-Base \\
\midrule
Standard & 46.31 & 75.65 & 69.13 \\
1-Layer & 50.39 & 75.55 & 68.81 \\
2-Layer & \textbf{51.90} & 75.99 & \textbf{71.78} \\
Full-Layer & 51.62 & \textbf{78.21} & 69.31 \\
\bottomrule
\end{tabular}%
}
\end{subtable}
\vspace{1em} % a bit of space between the two subtables
\begin{subtable}[t]{\linewidth}
\centering
\caption{Perturbation Coefficient Ablation (Learnable vs.\ Fixed DIPNet).}
\label{tab:learnable}
\resizebox{0.54\linewidth}{!}{%
\begin{tabular}{l c c c}
\toprule
Method & ViT-Tiny & ViT-Small & ViT-Base \\
\midrule
Standard & 46.31 & 75.65 & 69.13 \\
Fixed-0.5 & \textbf{51.91} & 77.02 & \textbf{69.95} \\
Fixed-1.0 & 48.94 & 75.97 & 67.23 \\
Fixed-1.2 & 47.44 & 75.24 & 65.93 \\
Fixed-Learned & 45.24 & 72.29 & 64.82 \\
Learnable & 51.62 & \textbf{78.21} & 69.31 \\
\bottomrule
\end{tabular}%
}
\end{subtable}
\end{table}
\subsubsection{Ablation Study on Layerwise Distributional Input Projection}
To assess how the depth of distributional input projections affects robustness, we compare four configurations of Layerwise DIPNet: taking the distributional input projection only at the first layer (``1-Layer''), at the first two layers (``2-Layer''), at every layer (``Full-Layer''), and the original model without any distributional projection (``Standard'').
As Table~\ref{tab:layerwise} shows, taking the distributional input projection at only a single layer (``1-Layer'') yields little to no generalization benefit—and can even lead to a slight drop in accuracy. In contrast, the ``2-Layer'' configuration delivers a clear generalization boost, outperforming Standard across all ViT scales, and on both ViT-Tiny and ViT-Base it even edges out the ``Full-Layer'' configuration—most notably pushing ViT-Base to 71.78~\%. Overall, taking the distributional input projection at multiple depths delivers substantial gains, although the optimal number of layers can vary by model scale.
\subsubsection{Ablation Study on Learnable Distributional Input Projection}
We then explore different strategies for the learnable parameter \texttt{coef}—which controls the magnitude of the distributional projection added at each layer—in DIPNet: fixing \texttt{coef} to a constant (0.5, 1.0 or 1.2) or the optimal value obtained from a learnable run (``Fixed-Learned''), or treating \texttt{coef} as a trainable parameter initialized randomly (``Learnable''). In the ``Fixed-Learned'' setting, we initialize \texttt{coef} to the value obtained from ``Learnable'' setting (1.4121 for ViT-Tiny, -1.4141 for ViT-Small, and 1.4141 for ViT-Base). As Table~\ref{tab:learnable} shows, fixing \texttt{coef} to a small constant produces an accuracy comparable to that of the ``Learnable'' setting, even though the ``Learnable'' setting itself starts from a random small value and grows to a larger optimum. However, choosing a larger fixed \texttt{coef} leads to a clear drop in performance, even when that constant remains below the final magnitude learned by the ``Learnable'' setting. This contrast highlights the advantage of a learnable \texttt{coef}, which can flexibly adjust perturbation strength to the ideal level for all input layers.
\section{Impact statement}\label{sec:impact}
This paper contributes foundational research in the areas of model generalization within the machine learning community. Our primary goal is to advance the practical methodologies to improve model generalization.
Given the scope of this research, we do not anticipate immediate ethical concerns or direct societal consequences. Therefore, we believe there are no specific ethical considerations or immediate societal impacts to be emphasized in the context of this work.
\begin{figure}[htb]
\centering
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_trainloss.png}
\caption{Training Loss}
\label{fig:train_loss}
\end{subfigure}
\hfill
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_valloss.png}
\caption{Validation Loss}
\label{fig:val_loss}
\end{subfigure}
\vspace{1em}
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_trainacc.png}
\caption{Training Accuracy}
\label{fig:train_acc}
\end{subfigure}
\hfill
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_valacc.png}
\caption{Validation Accuracy}
\label{fig:val_acc}
\end{subfigure}
\caption{Training and validation curves on ViT-Tiny without any attacks.}
\label{fig:vitt_curves}
\end{figure}
\begin{figure}[htb]
\centering
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitb_fgsm_trainloss.png}
\caption{Training Loss}
\label{fig:train_loss}
\end{subfigure}
\hfill
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitb_fgsm_valloss.png}
\caption{Validation Loss}
\label{fig:val_loss}
\end{subfigure}
\vspace{1em}
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitb_fgsm_trainacc.png}
\caption{Training Accuracy}
\label{fig:train_acc}
\end{subfigure}
\hfill
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitb_fgsm_valacc.png}
\caption{Validation Accuracy}
\label{fig:val_acc}
\end{subfigure}
\caption{Training and validation curves on ViT-Base under FGSM adversarial attacks.}
\label{fig:vitb_fgsm_curves}
\end{figure}
\begin{figure}[htb]
\centering
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_g_trainloss.png}
\caption{Training Loss}
\label{fig:train_loss}
\end{subfigure}
\hfill
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_g_valloss.png}
\caption{Validation Loss}
\label{fig:val_loss}
\end{subfigure}
\vspace{1em}
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_g_trainacc.png}
\caption{Training Accuracy}
\label{fig:train_acc}
\end{subfigure}
\hfill
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_g_valacc.png}
\caption{Validation Accuracy}
\label{fig:val_acc}
\end{subfigure}
\caption{Training and validation curves on ViT-Tiny under Gaussian noise.}
\label{fig:vitt_g_curves}
\end{figure}
\begin{figure}[htb]
\centering
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_f_trainloss.png}
\caption{Training Loss}
\label{fig:train_loss}
\end{subfigure}
\hfill
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_f_valloss.png}
\caption{Validation Loss}
\label{fig:val_loss}
\end{subfigure}
\vspace{1em}
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_f_trainacc.png}
\caption{Training Accuracy}
\label{fig:train_acc}
\end{subfigure}
\hfill
\begin{subfigure}[b]{0.45\textwidth}
\includegraphics[width=\linewidth]{pic/vit_pic/vitt_f_valacc.png}
\caption{Validation Accuracy}
\label{fig:val_acc}
\end{subfigure}
\caption{Training and validation curves on ViT-Tiny under FGSM adversarial attacks.}
\label{fig:vitt_fgsm_curves}
\end{figure}
\end{document}
\usepackage{amsfonts,amsmath,amssymb,amsthm,bm}
\usepackage{algorithm}
\usepackage{algorithmic}
\usepackage{mathrsfs}
%\usepackage[colorlinks,
% linkcolor=red,
% anchorcolor=blue,
% citecolor=blue
% ]{hyperref}
\usepackage{natbib}
\usepackage{tabularx}
\usepackage{multirow}
\usepackage{threeparttable}
\usepackage{makecell}
\usepackage{booktabs}
%\usepackage{fullpage}
\usepackage{appendix}
\usepackage{graphicx}
\usepackage{wrapfig}
\usepackage{subcaption}
\usepackage{caption}
\newtheorem{lemma}{Lemma}
\newtheorem{theorem}{Theorem}
\newtheorem{proposition}{Proposition}
\newtheorem{definition}{Definition}
\newtheorem{example}{Example}
\newtheorem{corollary}[theorem]{Corollary}
\newtheorem{remark}{Remark}
\newtheorem{condition}{Condition}
\ifx\assumption\undefined
\newtheorem{assumption}{Assumption}
\fi
\newcommand{\cF}{{\mathcal{F}}}
\newcommand{\cX}{{\mathcal{X}}}
\newcommand{\cA}{{\mathcal{A}}}
\newcommand{\cS}{{\mathcal{S}}}
\newcommand{\rR}{{\mathbb{R}}}
\newcommand{\E}{{\mathbb{E}}}
\newcommand{\V}{{\mathbb{V}}}
\newcommand{\Pb}{{\mathbb{P}}}
\newcommand{\Nb}{{\mathbb{N}}}
\newcommand{\reg}{\mathrm{Reg}}
\newcommand{\mdp}{\mathrm{MDP}}
\newcommand{\md}{\mathrm{d}}
\newcommand{\tv}{\mathrm{TV}}
\newcommand{\tPb}{\Tilde{\mathbb{P}}}
\newcommand{\tr}{\Tilde{r}}
\newcommand{\tV}{\Tilde{V}}
\newcommand{\tQ}{\Tilde{Q}}
\newcommand{\hQ}{\hat{Q}}
\newcommand{\hV}{\hat{V}}
\newcommand{\ks}{{k^{*}}}
\newcommand{\bias}{\mathbf{B}}
\newcommand{\variancen}{\mathbf{V}_{\operatorname{norm}}}
\newcommand{\biasn}{\mathbf{B}_{\operatorname{norm}}}
\renewcommand\bm[1]{{#1}}
\renewcommand\parallel{\|}
%\def\##1\#{\begin{align}#1\end{align}}
%\def\$#1\${\begin{align*}#1\end{align*}}
%\def\blue#1{\textcolor{blue}{#1}}
%\def\blue#1{\textcolor{black}{#1}}
\usepackage{xcolor} % colors
%\ifodd 0
\newcommand{\red}[1]{{\color{red}#1}}
\newcommand{\green}[1]{{\color{green}#1}}
\newcommand{\blue}[1]{{\color{blue}#1}}
%\newcommand\E{\mathbb{E}} % expected value
\newcommand\T{T} % transpose
\newcommand\truth{\theta} % true coefficients
\newcommand\est{\hat\theta} % arbitrary estimator
\newcommand\ridge{\hat\theta_\lambda} % ridge estimator
\newcommand\ols{\hat\theta_0} % least norm estimator
\newcommand\srisk{\mathcal{R}^{\operatorname{std}}} % standard risk
\newcommand\arisk{\mathcal{R}^{\operatorname{adv}}_{\budget}} % adversarial risk
\newcommand\budget{\alpha} % adversarial budget
\newcommand\sriskb{\mathcal{B}^{\operatorname{std}}}
\newcommand\sriskv{\mathcal{V}^{\operatorname{std}}}
\newcommand\normb{\mathcal{B}^{\operatorname{norm}}}
\newcommand\normv{\mathcal{V}^{\operatorname{norm}}}
\newcommand{\trace}{\operatorname{tr}}
\newcommand{\fracl}[2]{{#1}/{#2}}
\newcommand{\inprobto}{\overset{\text{pr.}}{\to}}
\begin{thebibliography}{}
\bibitem[Abbas et~al., 2022]{abbas2022sharp}
Abbas, M., Xiao, Q., Chen, L., Chen, P.-Y., and Chen, T. (2022).
\newblock Sharp-maml: Sharpness-aware model-agnostic meta learning.
\newblock In {\em International conference on machine learning}, pages 10--32. PMLR.
\bibitem[Bahri et~al., 2021]{bahri2021sharpness}
Bahri, D., Mobahi, H., and Tay, Y. (2021).
\newblock Sharpness-aware minimization improves language model generalization.
\newblock {\em arXiv preprint arXiv:2110.08529}.
\bibitem[Bousquet and Elisseeff, 2002]{bousquet2002stability}
Bousquet, O. and Elisseeff, A. (2002).
\newblock Stability and generalization.
\newblock {\em The Journal of Machine Learning Research}, 2:499--526.
\bibitem[Chen et~al., 2021]{chen2021vision}
Chen, X., Hsieh, C.-J., and Gong, B. (2021).
\newblock When vision transformers outperform resnets without pre-training or strong data augmentations.
\newblock {\em arXiv preprint arXiv:2106.01548}.
\bibitem[Cobbe et~al., 2021]{cobbe2021gsm8k}
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et~al. (2021).
\newblock Training verifiers to solve math word problems.
\newblock {\em arXiv preprint arXiv:2110.14168}.
\bibitem[Cohen et~al., 2019]{cohen2019certified}
Cohen, J., Rosenfeld, E., and Kolter, Z. (2019).
\newblock Certified adversarial robustness via randomized smoothing.
\newblock In {\em international conference on machine learning}, pages 1310--1320. PMLR.
\bibitem[Dhariwal and Nichol, 2021]{dhariwal2021diffusion}
Dhariwal, P. and Nichol, A. (2021).
\newblock Diffusion models beat gans on image synthesis.
\newblock {\em Advances in neural information processing systems}, 34:8780--8794.
\bibitem[Dobriban et~al., 2023]{dobriban2023provable}
Dobriban, E., Hassani, H., Hong, D., and Robey, A. (2023).
\newblock Provable tradeoffs in adversarially robust classification.
\newblock {\em IEEE Transactions on Information Theory}.
\bibitem[Donhauser et~al., 2021]{donhauser2021interpolation}
Donhauser, K., Tifrea, A., Aerni, M., Heckel, R., and Yang, F. (2021).
\newblock Interpolation can hurt robust generalization even when there is no noise.
\newblock {\em Advances in Neural Information Processing Systems}, 34:23465--23477.
\bibitem[Dosovitskiy et~al., 2020]{dosovitskiy2020vit}
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. (2020).
\newblock An image is worth 16x16 words: Transformers for image recognition at scale.
\newblock {\em arXiv preprint arXiv:2010.11929}.
\bibitem[Dubey et~al., 2024]{dubey2024llama3}
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et~al. (2024).
\newblock The llama 3 herd of models.
\newblock {\em arXiv preprint arXiv:2407.21783}.
\bibitem[Fanaee-T and Gama, 2014]{fanaee2014event}
Fanaee-T, H. and Gama, J. (2014).
\newblock Event labeling combining ensemble detectors and background knowledge.
\newblock {\em Progress in Artificial Intelligence}, 2:113--127.
\bibitem[Farabet and Warkentin, 2025]{gemma3}
Farabet, C. and Warkentin, T. (2025).
\newblock Introducing gemma 3: The most capable model you can run on a single gpu or tpu.
\newblock \url{https://blog.google/technology/developers/gemma-3/}.
\bibitem[Foret et~al., 2020]{foret2020sharpness}
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. (2020).
\newblock Sharpness-aware minimization for efficiently improving generalization.
\newblock {\em arXiv preprint arXiv:2010.01412}.
\bibitem[Gao et~al., 2024]{eval-harness}
Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., Le~Noac'h, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang, J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika, L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, A. (2024).
\newblock The language model evaluation harness.
\bibitem[Goodfellow et~al., 2014]{goodfellow2014explaining}
Goodfellow, I.~J., Shlens, J., and Szegedy, C. (2014).
\newblock Explaining and harnessing adversarial examples.
\newblock {\em arXiv preprint arXiv:1412.6572}.
\bibitem[Gowal et~al., 2020]{gowal2020uncovering}
Gowal, S., Qin, C., Uesato, J., Mann, T., and Kohli, P. (2020).
\newblock Uncovering the limits of adversarial training against norm-bounded adversarial examples.
\newblock {\em arXiv preprint arXiv:2010.03593}.
\bibitem[Hao and Zhang, 2024]{hao2024surprising}
Hao, Y. and Zhang, T. (2024).
\newblock The surprising harmfulness of benign overfitting for adversarial robustness.
\newblock {\em arXiv preprint arXiv:2401.12236}.
\bibitem[Ho et~al., 2020]{ho2020denoising}
Ho, J., Jain, A., and Abbeel, P. (2020).
\newblock Denoising diffusion probabilistic models.
\newblock {\em Advances in neural information processing systems}, 33:6840--6851.
\bibitem[Hochreiter and Schmidhuber, 1997]{hochreiter1997flat}
Hochreiter, S. and Schmidhuber, J. (1997).
\newblock Flat minima.
\newblock {\em Neural computation}, 9(1):1--42.
\bibitem[Izmailov et~al., 2018]{izmailov2018averaging}
Izmailov, P., Podoprikhin, D., Garipov, T., Vetrov, D., and Wilson, A.~G. (2018).
\newblock Averaging weights leads to wider optima and better generalization.
\newblock {\em arXiv preprint arXiv:1803.05407}.
\bibitem[Javanmard et~al., 2020]{javanmard2020precise}
Javanmard, A., Soltanolkotabi, M., and Hassani, H. (2020).
\newblock Precise tradeoffs in adversarial training for linear regression.
\newblock In {\em Conference on Learning Theory}, pages 2034--2078. PMLR.
\bibitem[Jiang et~al., 2021]{jiang2021assessing}
Jiang, Y., Nagarajan, V., Baek, C., and Kolter, J.~Z. (2021).
\newblock Assessing generalization of sgd via disagreement.
\newblock {\em arXiv preprint arXiv:2106.13799}.
\bibitem[Jiang et~al., 2019]{jiang2019fantastic}
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. (2019).
\newblock Fantastic generalization measures and where to find them.
\newblock {\em arXiv preprint arXiv:1912.02178}.
\bibitem[Johnson and Zhang, 2023]{johnson2023inconsistency}
Johnson, R. and Zhang, T. (2023).
\newblock Inconsistency, instability, and generalization gap of deep neural network training.
\newblock {\em arXiv preprint arXiv:2306.00169}.
\bibitem[Keskar et~al., 2016]{keskar2016large}
Keskar, N.~S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T.~P. (2016).
\newblock On large-batch training for deep learning: Generalization gap and sharp minima.
\newblock {\em arXiv preprint arXiv:1609.04836}.
\bibitem[Kirsch and Gal, 2022]{kirsch2022note}
Kirsch, A. and Gal, Y. (2022).
\newblock A note on" assessing generalization of sgd via disagreement".
\newblock {\em arXiv preprint arXiv:2202.01851}.
\bibitem[Krizhevsky, 2009]{krizhevsky2009learning}
Krizhevsky, A. (2009).
\newblock Learning multiple layers of features from tiny images.
\newblock Technical report, University of Toronto.
\newblock Technical Report.
\bibitem[Lecuyer et~al., 2019]{lecuyer2019certified}
Lecuyer, M., Atlidakis, V., Geambasu, R., Hsu, D., and Jana, S. (2019).
\newblock Certified robustness to adversarial examples with differential privacy.
\newblock In {\em 2019 IEEE symposium on security and privacy (SP)}, pages 656--672. IEEE.
\bibitem[Madry et~al., 2017]{madry2017towards}
Madry, A., Makelov, A., Schmidt, L., Tsipras, D., and Vladu, A. (2017).
\newblock Towards deep learning models resistant to adversarial attacks.
\newblock {\em arXiv preprint arXiv:1706.06083}.
\bibitem[Malinin et~al., 2022]{malinin2022shifts}
Malinin, A., Athanasopoulos, A., Barakovic, M., Cuadra, M.~B., Gales, M.~J., Granziera, C., Graziani, M., Kartashev, N., Kyriakopoulos, K., Lu, P.-J., et~al. (2022).
\newblock Shifts 2.0: Extending the dataset of real distributional shifts.
\newblock {\em arXiv preprint arXiv:2206.15407}.
\bibitem[Malinin et~al., 2021]{malinin2021shifts}
Malinin, A., Band, N., Chesnokov, G., Gal, Y., Gales, M.~J., Noskov, A., Ploskonosov, A., Prokhorenkova, L., Provilkov, I., Raina, V., et~al. (2021).
\newblock Shifts: A dataset of real distributional shift across multiple large-scale tasks.
\newblock {\em arXiv preprint arXiv:2107.07455}.
\bibitem[Moreno-Barea et~al., 2018]{moreno2018forward}
Moreno-Barea, F.~J., Strazzera, F., Jerez, J.~M., Urda, D., and Franco, L. (2018).
\newblock Forward noise adjustment scheme for data augmentation.
\newblock In {\em 2018 IEEE symposium series on computational intelligence (SSCI)}, pages 728--734. IEEE.
\bibitem[Nakkiran and Bansal, 2020]{nakkiran2020distributional}
Nakkiran, P. and Bansal, Y. (2020).
\newblock Distributional generalization: A new kind of generalization.
\newblock {\em arXiv preprint arXiv:2009.08092}.
\bibitem[Qin et~al., 2017]{qin2017dual}
Qin, Y., Song, D., Chen, H., Cheng, W., Jiang, G., and Cottrell, G. (2017).
\newblock A dual-stage attention-based recurrent neural network for time series prediction.
\newblock {\em arXiv preprint arXiv:1704.02971}.
\bibitem[Rombach et~al., 2022]{rombach2022high}
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. (2022).
\newblock High-resolution image synthesis with latent diffusion models.
\newblock In {\em Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, pages 10684--10695.
\bibitem[Saharia et~al., 2022]{saharia2022palette}
Saharia, C., Chan, W., Chang, H., Lee, C., Ho, J., Salimans, T., Fleet, D., and Norouzi, M. (2022).
\newblock Palette: Image-to-image diffusion models.
\newblock In {\em ACM SIGGRAPH 2022 conference proceedings}, pages 1--10.
\bibitem[Salman et~al., 2019]{salman2019provably}
Salman, H., Li, J., Razenshteyn, I., Zhang, P., Zhang, H., Bubeck, S., and Yang, G. (2019).
\newblock Provably robust deep learning via adversarially trained smoothed classifiers.
\newblock {\em Advances in neural information processing systems}, 32.
\bibitem[Shalev-Shwartz et~al., 2010]{shalev2010learnability}
Shalev-Shwartz, S., Shamir, O., Srebro, N., and Sridharan, K. (2010).
\newblock Learnability, stability and uniform convergence.
\newblock {\em The Journal of Machine Learning Research}, 11:2635--2670.
\bibitem[Shen and Meinshausen, 2023]{shen2023engression}
Shen, X. and Meinshausen, N. (2023).
\newblock Engression: Extrapolation for nonlinear regression?
\newblock {\em arXiv preprint arXiv:2307.00835}.
\bibitem[Song et~al., 2020]{song2020denoising}
Song, J., Meng, C., and Ermon, S. (2020).
\newblock Denoising diffusion implicit models.
\newblock {\em arXiv preprint arXiv:2010.02502}.
\bibitem[Tsipras et~al., 2018]{tsipras2018robustness}
Tsipras, D., Santurkar, S., Engstrom, L., Turner, A., and Madry, A. (2018).
\newblock Robustness may be at odds with accuracy.
\newblock {\em arXiv preprint arXiv:1805.12152}.
\bibitem[Vito, 2016]{misc_air_quality_360}
Vito, S. (2016).
\newblock {Air Quality}.
\newblock UCI Machine Learning Repository.
\newblock {DOI}: https://doi.org/10.24432/C59K5F.
\bibitem[Wang et~al., 2021]{wang2021convergence}
Wang, Y., Ma, X., Bailey, J., Yi, J., Zhou, B., and Gu, Q. (2021).
\newblock On the convergence and robustness of adversarial training.
\newblock {\em arXiv preprint arXiv:2112.08304}.
\bibitem[Yang et~al., 2024]{yang2024qwen2-5}
Yang, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Li, C., Liu, D., Huang, F., Wei, H., et~al. (2024).
\newblock Qwen2. 5 technical report.
\newblock {\em arXiv preprint arXiv:2412.15115}.
\bibitem[Yang et~al., 2020]{yang2020randomized}
Yang, G., Duan, T., Hu, J.~E., Salman, H., Razenshteyn, I., and Li, J. (2020).
\newblock Randomized smoothing of all shapes and sizes.
\newblock In {\em International Conference on Machine Learning}, pages 10693--10705. PMLR.
\bibitem[Yue et~al., 2023]{yue2023sharpness}
Yue, Y., Jiang, J., Ye, Z., Gao, N., Liu, Y., and Zhang, K. (2023).
\newblock Sharpness-aware minimization revisited: Weighted sharpness as a regularization term.
\newblock In {\em Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining}, pages 3185--3194.
\bibitem[Zhang et~al., 2019]{zhang2019theoretically}
Zhang, H., Yu, Y., Jiao, J., Xing, E., El~Ghaoui, L., and Jordan, M. (2019).
\newblock Theoretically principled trade-off between robustness and accuracy.
\newblock In {\em International conference on machine learning}, pages 7472--7482. PMLR.
\bibitem[Zhou et~al., 2021]{haoyietal-informer-2021}
Zhou, H., Zhang, S., Peng, J., Zhang, S., Li, J., Xiong, H., and Zhang, W. (2021).
\newblock Informer: Beyond efficient transformer for long sequence time-series forecasting.
\newblock In {\em The Thirty-Fifth {AAAI} Conference on Artificial Intelligence, {AAAI} 2021, Virtual Conference}, volume~35, pages 11106--11115. {AAAI} Press.
\bibitem[Zou et~al., 2021]{zou2021provable}
Zou, D., Frei, S., and Gu, Q. (2021).
\newblock Provable robustness of adversarial training for learning halfspaces with noise.
\newblock In {\em International Conference on Machine Learning}, pages 13002--13011. PMLR.
\end{thebibliography}