Sorting out Lipschitz function approximation

TL;DR

This paper introduces a new neural network architecture using GroupSort activation and norm constraints to effectively approximate Lipschitz functions, enhancing adversarial robustness and Wasserstein distance estimation.

Contribution

It proposes a novel architecture combining gradient norm preserving activation with norm-constrained weights, proving universality for Lipschitz function approximation.

Findings

Achieves tighter Wasserstein distance estimates than ReLU networks

Provides provable adversarial robustness with minimal accuracy loss

Demonstrates universality of the proposed Lipschitz approximator

Abstract

Training neural networks under a strict Lipschitz constraint is useful for provable adversarial robustness, generalization bounds, interpretable gradients, and Wasserstein distance estimation. By the composition property of Lipschitz functions, it suffices to ensure that each individual affine transformation or nonlinear activation is 1-Lipschitz. The challenge is to do this while maintaining the expressive power. We identify a necessary property for such an architecture: each of the layers must preserve the gradient norm during backpropagation. Based on this, we propose to combine a gradient norm preserving activation function, GroupSort, with norm-constrained weight matrices. We show that norm-constrained GroupSort architectures are universal Lipschitz function approximators. Empirically, we show that norm-constrained GroupSort networks achieve tighter estimates of Wasserstein distance than their ReLU counterparts and can achieve provable adversarial robustness guarantees with little cost to accuracy.