Towards Better Generalization via Distributional Input Projection Network

TL;DR

This paper introduces DIPNet, a novel neural network framework that projects inputs into learnable distributions at each layer, leading to smoother loss landscapes and improved generalization across diverse architectures and tasks.

Contribution

DIPNet is a new approach that enforces input smoothness via distributional projections, enhancing generalization in overparameterized models.

Findings

DIPNet reduces the Lipschitz constant of networks.

DIPNet improves test performance across various architectures.

DIPNet enhances robustness to adversarial and out-of-distribution inputs.

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.