PDE+: Enhancing Generalization via PDE with Adaptive Distributional Diffusion

TL;DR

This paper introduces PDE+ which leverages adaptive distributional diffusion within a PDE framework to enhance neural network generalization by smoothing solutions and covering unobserved data distributions.

Contribution

It establishes a PDE-based theoretical framework linking generalization to solution smoothness and proposes PDE+ to improve generalization through adaptive diffusion.

Findings

PDE+ outperforms state-of-the-art methods in various experiments.

The method improves model robustness to distribution shifts.

Adaptive diffusion enhances the smoothness of neural network solutions.

Abstract

The generalization of neural networks is a central challenge in machine learning, especially concerning the performance under distributions that differ from training ones. Current methods, mainly based on the data-driven paradigm such as data augmentation, adversarial training, and noise injection, may encounter limited generalization due to model non-smoothness. In this paper, we propose to investigate generalization from a Partial Differential Equation (PDE) perspective, aiming to enhance it directly through the underlying function of neural networks, rather than focusing on adjusting input data. Specifically, we first establish the connection between neural network generalization and the smoothness of the solution to a specific PDE, namely "transport equation". Building upon this, we propose a general framework that introduces adaptive distributional diffusion into transport equation to enhance the smoothness of its solution, thereby improving generalization. In the context of neural networks, we put this theoretical framework into practice as $\textbf{PDE+}$ ($\textbf{PDE}$ with $\textbf{A}$daptive $\textbf{D}$istributional $\textbf{D}$iffusion) which diffuses each sample into a distribution covering semantically similar inputs. This enables better coverage of potentially unobserved distributions in training, thus improving generalization beyond merely data-driven methods. The effectiveness of PDE+ is validated through extensive experimental settings, demonstrating its superior performance compared to SOTA methods.