Title: SORTING OUT LIPSCHITZ FUNCTION APPROXIMATION

Abstract: Training neural networks subject to a Lipschitz constraint is useful for generalization bounds, provable adversarial robustness, 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 function 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 Group-Sort 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.

:

Constraining the Lipschitz constant of a neural network ensures that a small change to the input can produce only a small change to the output. For classification, a small Lipschitz constant leads to better generalization (Sokolić et al., 2017) , improved adversarial robustness (Cisse et al., 2017; Tsuzuku et al., 2018) , and greater interpretability (Tsipras et al., 2018) . Additionally, the Wasserstein distance between two probability distributions can be expressed as a maximization problem over Lipschitz functions (Peyré & Cuturi, 2018) . But despite the wide-ranging applications, the question of how to approximate the class of Lipschitz functions with neural networks remains largely unanswered.

Existing approaches to enforce Lipschitz constraints broadly fall into two categories: regularization and architectural constraints. Regularization approaches such as double backprop (Drucker & Le Cun, 1992) or the gradient penalty (Gulrajani et al., 2017) perform well in practice, but do not provably enforce the Lipschitz constraint globally. On the other hand, norm-constrained architectures place limitations on the operator norm (such as the matrix spectral norm) of each layer's weight matrix (Cisse et al., 2017; Yoshida & Miyato, 2017) . These techniques provably satisfy the Lipschitz constraint, but this comes at a cost in expressive power. E.g., norm-constrained ReLU networks are provably unable to approximate simple functions such as absolute value (Huster et al., 2018) .

We first identify a simple property that expressive norm-constrained 1-Lipschitz architectures must satisfy: gradient norm preservation. Specifically, in order to represent a function with slope 1 almost everywhere, each layer must preserve the norm of the gradient during backpropagation. ReLU architectures satisfy this only when the activations are positive; empirically, this manifests during training of norm-constrained ReLU networks in that the activations are forced to be positive most of the time, reducing the network's capacity to represent nonlinear functions. We make use of an alternative activation function called GroupSort -a variant of which was proposed by Chernodub & Nowicki (2016) -which sorts groups of activations. GroupSort is both Lipschitz and gradient norm preserving. Using a variant of the Stone-Weierstrass theorem, we show that norm-constrained GroupSort networks are universal Lipschitz function approximators. While we focus our attention, both theoretically and empirically, on fully connected networks, the same general principles hold for convolutional networks where the techniques we introduce could be directly applied.

Empirically, we show that ReLU networks are unable to solve even the simplest Wasserstein distance estimation problems which GroupSort can solve completely. Moreover, we observe that norm-constrained ReLU networks must trade non-linear processing for gradient norm leading to less expressive networks. We also train classifiers with provable adversarial robustness guarantees and find that using GroupSort provides improved accuracy and robustness compared to ReLU. Across all of our experiments, we found that norm-constrained GroupSort architectures consistently outperformed their ReLU counterparts.

2. BACKGROUND:

Notation We will use x ∈ R in to denote the input vector to the neural network, y ∈ R out the output (or logits) of the neural network, n l the dimensionality of the l th hidden layer, W l ∈ R n l-1 ×n l and b l ∈ R n l the weight matrix and the bias of the l th layer. We will denote the pre-activations in layer l with z l and activations with h l . The number of layers in the network will be L with y = z L . We will use φ to denote the activation function used in the neural network. The computation performed by layer l of the network will be:

z l = W l h l-1 + b l h l = φ(z l )

Network Jacobian Using the chain rule, the Jacobian of a neural network can be expanded as follows:

∂y ∂x = ∂z L ∂h L-1 ∂h L-1 ∂z L-1 . . . ∂z 2 ∂h 1 ∂h 1 ∂z 1 ∂z 1 ∂x = W L φ (z L-1 ) . . . W 2 φ (z 1 )W 1 (1)

2.1 LIPSCHITZ FUNCTIONS Given two metric spaces X and Y, a function f : X → Y is Lipschitz continuous if there exists K ∈ R such that for all x 1 and x 2 in X ,

d Y (f (x 1 ), f (x 2 )) ≤ Kd X (x 1 , x 2 )

where d X and d Y are metrics (such as Euclidean distance) on X and Y respectively. In this work, when we refer to the Lipschitz constant we are referring to the smallest such K for which the above holds under a given d X and d Y . Unless otherwise specified, we take X = R n and Y = R m throughout. If the Lipschitz constant of a function is K, it is called a K-Lipschitz function.

Equivalently, if the function is everywhere differentiable then its Lipschitz constant is bounded by the operator norm of its Jacobian. Throughout this work, we make use of the following definition, Definition 1. Given a metric space (X, d X ) where d X denotes the metric on X, we write C L (X, R) to denote the space of all 1-Lipschitz functions mapping X to R (with respect to the L p metric).

2.2. LIPSCHITZ-CONSTRAINED NEURAL NETWORKS:

As 1-Lipschitz functions are closed under composition, to build a 1-Lipschitz neural network it suffices to compose 1-Lipschitz affine transformations and activation functions.

1-Lipschitz Linear Transformations: Ensuring that each linear map is 1-Lipschitz is equivalent to ensuring that ||Wx|| p ≤ ||x|| p for any x; i.e. constraining the matrix p-norm, ||W|| p = sup ||x||p=1 ||Wx|| p , to be at most 1. Important examples of matrix p-norms include the matrix 2-norm, which is the largest singular value, and the matrix ∞-norm, which can be expressed as:

||W|| ∞ = max 1≤i≤m m j=1 |w ij |.

Similarly, we may also define the mixed matrix norm, given by ||W|| p,q = sup ||x||p=1 ||Wx|| q . Enforcing matrix norm constraints naively may be computationally expensive. Fortunately, techniques exist to efficiently ensure that ||W || p = 1 when p = 2 or p = ∞. We discuss these in Section 4.2.

1-Lipschitz Activation Functions::

Most commonly used activation functions (such as ReLU (Krizhevsky et al., 2012 ), sigmoid, tanh, maxout (Goodfellow et al., 2013) ) are 1-Lipschitz, if they are scaled appropriately.

2.3. APPLICATIONS OF LIPSCHITZ NETWORKS:

Wasserstein Distance Estimation Wasserstein-1 distance (also called Earth Mover Distance) is a principled distance metric between two probability distributions and has found many applications in machine learning in recent years (Peyré & Cuturi, 2018; Genevay et al., 2017) . Using Kantorovich-Rubinstein duality (Villani, 2008) , one can recast the Wasserstein distance estimation problem as a concave maximization problem, defined over 1-Lipschitz functions:

EQUATION

Since this dual objective resembles the discriminator objective for generative adversarial networks (GANs), Arjovsky et al. (2017) proposed the Wasserstein GAN architecture, which uses a neural net architecture to approximate the space of Lipschitz functions.

Adversarial Robustness:

Adversarial examples are inputs to a machine learning system which have been designed to force undesirable behaviour (Szegedy et al., 2013; Goodfellow et al., 2014) . Formally, given a classifier f and a data point x, we write an adversarial example as x adv = x + δ such that f (x adv ) = f (x) and δ is small. A small Lipschitz constant can guarantee a lower bound on the size of δ (Tsuzuku et al., 2018) and thus provide robustness guarantees. However, existing approaches have both practical and theoretical limitations (Huster et al., 2018) .

3. GRADIENT NORM PRESERVATION:

When backpropagating through a norm-constrained 1-Lipschitz network, the gradient norm is nonincreasing as it is processed by each layer. This simple fact leads to interesting consequences when we attempt to represent (scalar-valued) functions whose input-output gradient has norm 1 almost everywhere. (Such functions are relevant to Wasserstein distance estimation, where an optimal dual solution always has this property (Gulrajani et al., 2017) .) To ensure the input-output gradient norm is 1, the gradient norm must be preserved by each layer in the network during backpropagation.

Unfortunately, norm-constrained networks with common activations are unable to achieve this. Theorem 1. Consider a neural network, f : R n → R, built with matrix 2-norm constrained weights (||W|| 2 ≤ 1) and 1-Lipschitz, element-wise, monotonically increasing activation functions.

If ||∇f (x)|| 2 = 1 almost everywhere, then f is linear.

As a special case, Theorem 1 shows that no 2-norm-constrained neural network with ReLU (or sigmoid, tanh, etc.) activations can represent the absolute value function. A full proof can be found in Appendix C. Informally, for ReLU layers, the gradient norm can only be preserved if every activation is positive (with the exception of units which don't affect the network's output). But as this holds for almost all inputs, the network's input-output mapping must be linear.

This tension between gradient norm and nonlinear processing is also observed empirically. Figure 5 compares the activation statistics for MNIST classification networks with ReLU activations, with and without matrix norm constraints on the weights. For the network with smallest Lipschitz constant, around 10% of the units are "undead", or always active (and hence do not contribute any nonlinear processing). This suggests that the network is sacrificing nonlinear capacity in order to maintain adequate gradient norm.

Another useful consequence of gradient norm preservation is that we may restrict all of the weight matrices to have singular values of 1: Theorem 2. Consider a neural network, f : R n → R, built with matrix 2-norm constrained weights and with ||∇f (x)|| 2 = 1 almost everywhere. Then, without changing the computed function, each weight matrix W ∈ R m×k can be replaced with a matrix W whose singular values all equal 1.

Note that the condition of singular values equaling 1 is equivalent to the following: when m > k, the columns of W are orthonormal; when m < k, the rows of W are orthonormal; and when m = k, W is orthogonal. For the remainder of this paper, we abuse terminology slightly and refer to such matrices as orthonormal. The proof of Theorem 2 is given in Appendix C. With these two results in place, we restrict our search for expressive Lipschitz networks to those that contain orthonormal weight matrices (those with singular values all equal to 1) and activations which preserve the gradient norm during backpropagation.

4. METHODS:

We begin by observing that if we can learn any 1-Lipschitz function with a neural network then we can trivially extend this to K-Lipschitz functions by scaling the output by K. With this in mind, we focus on designing 1-Lipschitz network architectures with respect to the L 2 and L ∞ metrics by requiring each layer to be 1-Lipschitz.

Figure 1 : GroupSort activation with a grouping size of 5.

4.1. GRADIENT NORM PRESERVING ACTIVATION FUNCTIONS:

As discussed in Section 3, commonly used activation functions such as ReLU are not gradient norm preserving. To achieve norm preservation, we use a general purpose 1-Lipschitz activation function which we call GroupSort. This activation function takes a column vector x ∈ R n , separates the elements into g groups, sorts each group into ascending order, and outputs the combined "group sorted" vector. This is shown graphically in Figure 1 .

Properties of GroupSort GroupSort is a Lipschitz operation. Furthermore, it is norm preserving: its Jacobian is a permutation matrix, and permutation matrices preserve every vector p-norm. Note also that GroupSort is homogeneous, i.e. GroupSort(αx) = αGroupSort(x), since the sorting order of the elements is invariant to scaling.

Varying the Grouping Size When we pick a grouping size of 2 for GroupSort, we call the operation MaxMin. This is equivalent to the Orthogonal Permutation Linear Unit (OPLU) activation (Chernodub & Nowicki, 2016) , which was also motivated based on gradient norm preservation.

When sorting the entire input vector, we call the operation FullSort. GroupSort, MaxMin, and Full-Sort are equally expressive, i.e. they can all be reduced to each other, such that the reduction obeys the norm constraint on the weights (for any matrix p-norm). We present the details in Appendix A. Compared to MaxMin, FullSort is able to represent certain functions more compactly, but we find that it is typically more difficult to train via stochastic gradient descent.

Representing other activations Under the matrix 2-norm constraint, MaxMin can be seen as equivalent to absolute value. We describe exactly how these activation functions can be transformed into each other in Appendix A. Applying absolute value to the activations has the effect of folding the space on each of the coordinate axes. Hence, a rigid linear transformation, followed by absolute value, followed by another rigid linear transformation, can implement folding along an arbitrary hyperplane. This gives an interesting interpretation of how MaxMin networks can represent certain functions by way of implementing absolute value; an example is shown in Figure 10 in Appendix A. Montufar et al. (2014) provide an in-depth analysis of the expressivity of neural networks built with activations that can perform folding.

Without norm constraints, GroupSort can recover many other common activation functions. For example, ReLU, Leaky ReLU, concatenated ReLU (Shang et al., 2016) , and maxout. Details can be found in Appendix A.

4.2. NORM-CONSTRAINED LINEAR MAPS:

We discuss how to practically enforce the 1-Lipschitz constraint on linear layers for 2-and ∞-norms.

4.2.1. ENFORCING ||W ||:

2 = 1 WHILE PRESERVING GRADIENT NORM

Several methods have been proposed to enforce matrix 2-norm constraints during training (Cisse et al., 2017; Yoshida & Miyato, 2017) . However, in the interest of preserving the gradient norm, we go a step further and enforce orthonormality of the weight matrices in each layer. This is a stronger condition, in that we require that all singular values be exactly 1, rather than bounded by 1.

We make use of an algorithm first introduced by Björck & Bowie (1971) , which we refer to as Björck Orthonormalization (or simply Björck). Given a matrix, this algorithm finds the closest orthonormal matrix through an iterative application of the Taylor expansion of the polar decomposition. Given an input matrix A 0 = A, the algorithm computes,

EQUATION

where

Q k = I -A T k A k .

Importantly, this algorithm is fully differentiable and thus has a pullback operator for the Stiefel manifold (Absil et al., 2009) allowing us to optimize over orthonormal matrices directly. A larger choice of p adds more computation but gives a closer approximation for each iteration. In practice, we found that we could use p = 1 with 2-3 iterations per forward pass and increase this to 15 or more iterations at the end of training to ensure a tightly enforced Lipschitz constraint. We discuss additional details of this algorithm including comparisons to Parseval networks (Cisse et al., 2017) and spectral normalization (Miyato et al., 2018) in Appendix B.

Note that while we focus on fully connected layers, the same general principles apply to convolutions. Convolutions can be unfolded and represented as a linear transformation. Up to constant rescaling, the spectral norm of the filter then bounds the spectral norm of the unfolded operation. We do not devote space to computing these constants but instead point readers to other resources which address this question (Gouk et al., 2018; Cisse et al., 2017; Sedghi et al., 2018) .

4.2.2. ENFORCING ||W || ∞ = 1:

Due to its simplicity and suitability for a GPU implementation, we use Algorithm 1 from Condat (2016) to project the weight matrices onto the L ∞ ball in all of our experiments. Other more sophisticated methods can be found in Condat (2016) .

4.3. PROVABLE ADVERSARIAL ROBUSTNESS:

A small Lipschitz constant limits the change in network output under small adversarial perturbations. As explored by Tsuzuku et al. (2018) , we can guarantee adversarial robustness at a point x by considering the margin about that point divided by the Lipschitz constant. Formally, given a network with Lipschitz constant K (with respect to the L ∞ metric) and an input x with corresponding class t that produces logits y, we define its margin by

EQUATION

If M(x) > K /2, then the network is robust to all adversarial perturbations δ with ||δ|| ∞ < , at x. In this work we train networks with ∞-norm constraints on their weights using a multi-class hinge loss:

EQUATION

where κ controls the margin enforcement and depends on the Lipschitz constant and desired perturbation tolerance (e.g. κ = 0.3 × K).

5. RELATED WORK:

Several methods have been proposed to train Lipschitz neural networks (Cisse et al., 2017; Yoshida & Miyato, 2017; Miyato et al., 2018; Gouk et al., 2018) . Cisse et al. (2017) regularize the weights of neural networks to obey an orthonormality constraint and utilize Lipschitz activation functions. In fact, the corresponding update to the weights due to this regularization term can be seen as one step of the Björck orthonormalization scheme (Equation 3). Another approach, spectral normalization (Miyato et al., 2018) , employs an efficient implementation of power iteration to rescale each weight by its spectral norm. We compare these methods to Björck orthonormalization in Appendix B.

Other researchers (Arjovsky et al., 2016; Wisdom et al., 2016; Sun et al., 2017) have explicitly parameterized square orthogonal weight matrices using, for example, Householder transformations (Householder, 1958) .

Other regularization techniques penalize the network Jacobian, thereby constraining the Lipschitz constant locally around the data (Gulrajani et al., 2017; Drucker & Le Cun, 1992; Sokolić et al., 2017) . While these methods have the advantage that it is typically easy to train neural networks under such penalties, they do not provably enforce a Lipschitz constraint. Gulrajani et al. (2017) apply the gradient penalty at randomly sampled points between two distributions, but as shown by Gemici et al. (2018) , this is often sub-optimal in the context of Wasserstein Distance estimation.

The Lipschitz constant of a neural network has been connected theoretically and empirically to its generalization performance (Bartlett, 1998; Bartlett et al., 2017; Neyshabur et al., 2017; 2018; Sokolić et al., 2017) . Neyshabur et al. (2018) show that if the network Lipschitz constant is small then a non-vacuous bound on the generalization error can be derived. Small Lipschitz constants have also been linked to adversarial robustness (Tsuzuku et al., 2018; Cisse et al., 2017) . In fact, adversarial training can be viewed as approximate gradient regularization (Miyato et al., 2017; Simon-Gabriel et al., 2018) which makes the function Lipschitz locally around the training data. Lipschitz constants have been used to provide provable adversarial robustness guarantees. Tsuzuku et al. (2018) manually enforce a margin depending on an approximation of the upper bound on the Lipschitz constant which in turn guarantees adversarial robustness. In this work we also explore provable adversarial robustness through margin training but do so with a network whose Lipschitz constant is known and globally enforced.

Classic neural network universality results use constructions which violate the norm-constraints needed for Lipschitz guarantees (Cybenko, 1989; Hornik, 1991) . Huster et al. (2018) explored universal approximation properties of ∞-norm-constrained networks and proved that ReLU activations cannot be used to approximate the absolute value function. In this work we also show that many activations, including ReLU, are deficient with 2-norm constraints. However, we prove that Lipschitz functions can be universally approximated if the correct activation function is used.

6. UNIVERSAL APPROXIMATION OF LIPSCHITZ FUNCTIONS:

Universal approximation results for general continuous functions do not directly apply to Lipschitz networks as the constructions typically involve huge Lipschitz constants. Moreover, Huster et al. (2018) showed that it is impossible to approximate even the absolute value function with ∞-normconstrained ReLU networks. In this section, we present theoretical guarantees on the approximation of Lipschitz functions with norm-constrained neural networks. To our knowledge, this is the first universal Lipschitz function approximation result for norm-constrained neural networks.

We will first prove a variant of the Stone-Weierstrass Theorem which gives a simple criterion for universality. (A similar result is presented in Lemma 4.1 in Yaacov (2010).) We then construct a class of networks with the GroupSort activation which satisfy this criterion. We now proceed with the formal statements. Definition 2. We say that a set of functions, L, is a lattice if for any f, g ∈ L we have max(f, g) ∈ L and min(f, g) ∈ L (where max and min are defined pointwise).

Lemma 1. (Restricted Stone-Weierstrass Theorem) Suppose that (X, d X ) is a compact metric space with at least two points and L is a lattice in C L (X, R) with the property that for any two distinct elements x, y ∈ X and any two real numbers a and b such that |a -b| ≤ d X (x, y) there exists a function f ∈ L such that f (x) = a and f (y

) = b. Then L is dense in C L (X, R).

Remark. We could replace | • | with any metric on R.

The full proof of Lemma 1 is presented in the appendix. Note that Lemma 1 says that A is a universal approximator for 1-Lipschitz functions if and only if A is a lattice that separates points.

Using Lemma 1, we can derive the second of our key results. Norm-constrained networks with GroupSort activations are able to approximate any Lipschitz function in L p distance. Theorem 3. (Universal Approximation with Lipschitz Networks) Let LN p denote the class of fullyconnected neural networks whose first weight matrix satisfies ||W 1 || p,∞ = 1, all other weight matrices satisfy ||W|| ∞ = 1, and MaxMin activations. Let X be a closed and bounded subset of R n endowed with the L p metric. Then the closure of

LN p is dense in C L (X, R). Proof. (Sketch) Observe first that LN p ⊂ C L (X, R).

By Lemma 1, it is sufficient to show that LN p is closed under max and min and has the point separation property. For the latter, note that given x, y ∈ X and a, b ∈ R with |a -b| ≤ ||x -y|| p , we can fit a line with a single layer network, f , satisfying the 1-Lipschitz constraint with f (x) = a and f (x) = b.

Now consider f and g in LN p . For simplicity, here assume that they have the same number of layers. We can construct h ∈ LN ∞ by taking the weight matrix of the first layer to be the weight matrices of the first layer in f and g vertically concatenated. For the following layers, instead of vertically stacking, we build a block diagonal matrix from the weights of f and g. This network is in LN p and the final layer of the network outputs [f (x), g(x)]. We then apply the GroupSort activation to get [max(f, g)(x), min(f, g)(x)] and finally take the dot product with [1, 0] or [0, 1] to get the max or min respectively.

We refer readers to Appendix D for the formal proof of Theorem 3 and a diagram of the constructed network in Figure 14 . One special case of Theorem 3 is for 1-Lipschitz functions in L ∞ norm, While these constructions rely on the matrix ∞-norm of the weight matrices being constrained, we find in practice that constraining the matrix 2-norm makes the networks easier to We propose a simple yet effective method to quantify how expressive different Lipschitz architectures are. We first carefully pick pairs of probability distributions whose Wasserstein Distance and (unique) optimal dual surfaces can be computed analytically. Then, we train neural networks to optimize the dual Wasserstein distance objective (Equation 2) using samples from these distributions and compare the estimated Wasserstein distance and learned dual surfaces to the optimal, analytically computed ones. Expressiveness is measured by how closely the neural network can estimate the correct Wasserstein distance. For 1D and 2D problems, the learned dual surfaces can also be visualized, making it possible to inspect failure modes of non-expressive architectures.

In the following experiments, we trained networks to approximate the absolute value function, multiple two dimensional circular cones and single high dimensional circular cones. Appendix G.1 describes how pairs of probability distributions can be picked which have these optimal dual surfaces, and a Wasserstein distance of precisely 1. In all of the experiments in this section, we use Björck orthonormalization to enforce the 2-norm constraints on the weights.

Approximating absolute value function: Figure 2 shows the dual surfaces approximated by Lipschitz-constrained networks with various activation functions. The optimal dual surface is the absolute value. It can be seen that non-GNP activation functions are incapable of approximating this rather trivial Lipschitz function. While we observed that increasing the network depth helps ReLU and MaxOut activations (Table 1 ), the representational bottleneck showcased in Figure 2 leads to more severe limitations as the problem dimensionality increases.

Approximating multiple 2D cones: Figure 3 shows the dual surfaces approximated by neural networks using various activations. The optimal dual surface is three consecutive circular cones with a gradient of 1 everywhere. Here we observed an even more serious pathology with non-GNP activations: by attempting to increase the slope, the non-GNP networks may distort the shape of the dual surface. When training WGAN critics, this problem cannot be fixed by increasing the Lipschitz constant, since optimal critics for different Lipschitz constants are equivalent up to scaling.

Approximating high dimensional circular cones: We evaluated the performance of architectures built with different activation functions for higher dimensional inputs, on the task of approximating high dimensional circular cones which have a gradient of 1 everywhere. As shown in Table 1 , this leads to significant drops in the Wasserstein dual objective for Lipschitz networks built with non-GNP activations, and increasing the depth of the networks only slightly improves the situation. We also observed that while the MaxMin activation performs significantly better, it also needs large depth in order to learn the optimal solution. Surprisingly, the FullSort network has no difficulty approximating high dimensional circular cones, even with only two hidden layers. The ReLU network was not able to make use of the capacity afforded to it and the observed Lipschitz constant was actually closer to 8 than 10. In Appendix F.3 we show the full singular value distribution which suggests that 2-norm-constrained MaxMin networks can achieve dynamical isometry (Pennington et al., 2017) throughout training.

We studied the activation statistics of ReLU networks trained to classify MNIST digits with and without 2-norm constraints in Figure 5 . Given a threshold value, τ ∈ [0, 1], we computed the proportion of activations throughout the network which are positive at least as often as τ over the training data distribution. Without a Lipschitz constraint, the activation statistics were very sparse, with almost no units active when τ > 0.4, even when using dropout (Srivastava et al., 2014) . When the Lipschitz constraint was enforced the activations were much less sparse with smaller Lipschitz constants amplifying the effect. In the worst case, about 10% of units were "undead", or active all of the time, and hence did not contribute any nonlinear processing. It's not clear what effect this has on the network's representational capacity, but such a dramatic change in the network's activation statistics suggests that it made significant compromises in order to maintain adequate gradient norm.

7.2. WASSERSTEIN DISTANCE ESTIMATION:

We turn our attention to using norm-constrained GroupSort networks to estimate the Wasserstein distance between the generator distribution of a GAN and the empirical distribution of the data it was trained on. We note that optimal surfaces under the dual Wasserstein objective have a gradient norm of 1 almost everywhere (Corollary 1 in Gemici et al. ( 2018)). Hence, the gradient norm preservation properties discussed in Section 3 are critical. Appendix G.2 contains details on the experiments described in this section.

7.2.1. LOWER BOUNDS ON MNIST AND CIFAR-10 GANS:

In this experiment, we first trained a GAN variant on MNIST and CIFAR-10 datasets and then froze the weights of the generator. Using samples from the generator and original data distribution, we trained independent 1-Lipschitz neural networks to compute the Wasserstein distance between the empirical data distribution and the generator distribution. As can be seen in Table 2 , using normpreserving activation functions helps achieve a tighter lower bound on the Wasserstein distance for both MNIST and CIFAR-10 generators.

Training WGANs We were also able to train WGANs using our proposed 1-Lipschitz activations and linear transformations. We borrowed the discriminator and generator architectures directly from Chen et al. ( 2016), but switched the ReLU activations with MaxMin and replaced the standard convolutional and fully connected layers with their Björck counterparts. We also dropped the batch normalization layers, as these would violate the Lipschitz constraint. Figure 6 shows MNIST and CIFAR-10 samples generated using our WGAN variant. We leave further investigation of the WGANs built with our techniques to a future study.

7.3. ROBUSTNESS AND INTERPRETABILITY OF LIPSCHITZ NETWORKS:

We explored the robustness of Lipschitz neural networks trained on MNIST to adversarial perturbations measured with L ∞ distance. When training the networks we enforced an L ∞ constraint on the weights and used the multi-class hinge loss from Equation 5. We found this to be more effective than the manual margin training used by Tsuzuku et al. (2018) . We trained all networks with a Lipschitz constant of K = 1000 and chose the margin κ = Ka where a was 0.1 or 0.3. Notably, this technique provides margin-based provable robustness guarantees as described in Section 4.3. We then attacked these models using the FGS and PGD methods (Szegedy et al., 2013; Madry et al., 2017) under the CW loss (Carlini & Wagner, 2016) . The results are presented in Table 3 and Figure 8 . The Lipschitz networks with MaxMin activations were able to achieve better clean accuracy and larger margins than their ReLU counterparts which led to considerably improved adversarial robustness.

Figure 6 : Samples from WGANs whose critic architectures were built using GNP atomic units. With the strictly enforced Lipschitz constant, we can compute theoretical lower bounds on the accuracy against adversaries with a maximum perturbation strength . In Figure 9 , we show this lower bound for each of the models previously studied. This is computed by finding the proportion of data points which violate the margin by at least K . Note that at the computed threshold, the model has low confidence in the adversarial example. An even larger perturbation would be required to induce confident misclassification. 

8. CONCLUSION:

We have identified gradient norm preservation as a critical component of Lipschitz network design and showed that failure to achieve this leads to less expressive networks. By combining the Group-Sort activation function and orthonormal weight matrices, we presented a class of neural networks which are provably 1-Lipschitz and can approximate any 1-Lipschitz function arbitrarily well. Empirically, we showed that our GroupSort networks are more expressive than existing architectures and can be used to achieve better estimates of Wasserstein distance and provable adversarial robustness guarantees.

Appendices:

A GROUPSORT ACTIVATION FullSort and MaxMin FullSort can implement MaxMin by simply "chunking" the biases in pairs. To be more precise, let x max = sup x∈X ||x|| ∞ where X represents the domain, and b

= [b 1 , b 2 , ..., b n ] T where x max < b 1 = b 2 b 3 = b 4 • • • b n-1 = b n (

denotes differing by at least x max ). We can write:

MaxMin(x) = FullSort(Ix + b) -b,

where I denotes the identity matrix. Similarly, FullSort can be represented using a series of MaxMin layers that implement BubbleSort; note that this construction obeys any matrix p-norm constraint since it can be implemented using only permutation matrices for the weights.

MaxMin and absolute value MaxMin and absolute value can each represent eachother under 2-norm-constrained weights. The two operations are reduced to each other as follows:

max(x) min(y) = 1 √ 2 1 √ 2 1 √ 2 -1 √ 2 abs( 1 √ 2 1 √ 2 1 √ 2 -1 √ 2 x y + B 0 ) - √ 2B 0 (6) abs(x) = 1 √ 2 -1 √ 2

MaxMin(

EQUATION

In Equation 6, the value of B is chosen such that 2x + √ 2B > 0 for all x in the domain. Note that all the matrices in these constructions satisfy the matrix 2-norm constraint.

Figure 10 : A rigid linear transformation, followed by absolute value, followed by another rigid linear transformation, can implement folding along an arbitrary hyperplane. Here is an example where the network represents a function consisting of a pair of square pyramids by folding the space three times, until the function is representable as a linear function of the top layer activations.

GroupSort and other activations Here we show that GroupSort can recover ReLU, maxout, and concatenated ReLU activation functions. We first show that MaxMin can recover ReLU and concatenated ReLU. Note that,

EQUATION

Thus, by adding 0 elements to the pre-activations and then applying another linear transformation after MaxMin we can output either ReLU or concatenated ReLU. If instead of adding 0 to the preactivations we added ax we could recover Leaky ReLU by using a linear transformation to select max(x, ax).

To recover maxout with groups of size k, we perform GroupSort with groups of size k and use the next linear transformation to select the first element of each group after sorting (corresponding to the max).

B IMPLEMENTING NORM CONSTRAINTS:

When implementing the norm constraints it is possible to project the weight matrices after each gradient descent step, or during the forward pass (if the projection is differentiable). For the Björck algorithm we utilize the latter while Parseval networks use the Björck algorithm after each gradient descent step. In any case, once training has completed we can project the weights to enforce the norm constraint and use these as our fixed weights at test time -removing the computational overhead required during training.

B.1 COMPARING BJ ÖRCK AND PARSEVAL:

In Cisse et al. (2017) , the authors motivate an update to the weight matrices by considering the gradient of a regularization term, β 2 ||W T W -I|| 2 F . By this gradient from the weight matrices they push them closer to the Stiefel manifold. The final update is given by,

W ← W (I + β) -βW W T W (9)

Note that when β = 0.5 this update is exactly the first order (p = 1) update from Equation 3, with a single iteration. Compared to our approach, the key difference in Parseval networks is that the weight matrix update is applied after the primary gradient update. For our approach, we utilize the algorithm in Equation 3 during the network forward pass to optimize directly on the Stiefel manifold. This is more expensive but lets us ensure that the weight matrices are close to orthonormal throughout training.

Choice of β:

We can relate the first order Björck algorithm to the Parseval update by setting β = 0.5. However, in practice Parseval networks are trained with very small choices of β, for example β = 0.0003. As expected, when β is small the algorithm still converges to an orthonormal matrix but much more slowly. Figure 11 shows the maximum and minimum singular values of matrices which have undergone 50 iterations of the first order Björck scheme for varying choices of β < 0.5. When β is much smaller than 0.5 the matrices may be far from orthonormal. We also show how the maximum and minimum singular values vary over the number of iterations when β = 0.0003 (a common choice for Parseval networks) in Figure 12 . This has practical implications for Parseval training, particularly when using early stopping, as the weight matrices may be far from orthonormal if the gradients are relatively large compared to the update produced by the Björck algorithm. We observed this effect empirically in our MNIST classification experiments but found that Parseval networks were still able to achieve a meaningful regularization effect.

B.2 COMPARING BJ ÖRCK AND SPECTRAL NORMALIZATION:

Spectral Normalization (Miyato et al., 2018) enforces the largest singular value of each weight matrix to be less than 1 by estimating the largest singular value and left/right singular vectors using power iteration, and normalizing the weight matrix using these during each forward pass. While this constraint does allow all singular values of the weight matrix to be 1, we have found that this rarely happens in practice. Hence, enforcing the 1-Lipschitz constraint via spectral normalization doesn't guarantee gradient norm preservation.

We demonstrate the practical consequences of the inability of spectral normalization to preserve gradient norm on the task of approximating high dimensional cones. In order to quantify approximation performance, we carefully pick two n dimensional probability distributions such that 1) The Wasserstein Distance between them is exactly 1 and 2) the optimal dual surface consists of an n -1 dimensional cones with a gradient of 1 everywhere, embedded in n dimensions. We trained 1-Lipschitz constrained neural networks to optimize the dual Wasserstein objective in 2 and checked how well the architecture of choice is able to approximate the optimal dual surface, measured by the Wasserstein Distance they estimate. Please refer to Section 7.1.1 for more experiments in this flavor and Appendix G.1 for how these two probability distributions are picked.

Figure 13 shows that neural networks trained with Björck orthonormalization not only are able to approximate high dimensional cones better than spectral normalization, but also converge much faster in terms of training iterations. The gap between these methods gets much more significant as the problem dimensionality increases. In this experiment, each network consisted of 3 hidden layers Figure 13 : Comparing the performance of 1-Lipschitz neural nets using Björck orthonormalization and spectral normalization to enforce the 2-norm constraint on the high dimensional cone fitting task (Section 7.1.1). Note that networks using Björck orthonormalization both converge faster and achieve higher final approximation accuracies, measured by the estimated Wasserstein Distance.

with 512 hidden units per layer, and was trained with the Adam optimizer (Kingma & Ba, 2014) with its default hyperparameters. Tuned learning rates of 0.01 for Björck and 0.0033 for spectral normalization were used.

B.3 SUFFICIENT CONDITION FOR CONVERGENCE OF BJ ÖRCK ORTHONORMALIZATION:

The Björck orthonormalization can be shown to always converge as long as the condition ||W T W-I|| 2 < 1 is satisfied (Hasenclever et al.) . When viewed in conjunction with the fact that the output of this procedure is scale-invariant (BJORCK(αW) = αBJORCK(W)) (Björck & Bowie, 1971) , the aforementioned sufficient condition can be implemented by simply scaling the weight matrix so that all of its singular values are smaller than or equal to 1 before orthonormalization.

A scaling factor can be computed efficiently by considering the following matrix norm inequalities:

σ max ≤ √ m * n||W|| max (10) σ max ≤ √ n||W|| 1 (11) σ max ≤ √ m||W|| ∞ (12)

Above, σ max corresponds to the largest singular value of the matrix and m and n stand for the number of rows and columns respectively. Note that computing the quantities on the right hand side of the inequalities involves at most summing over the rows or columns of the weight matrix, which is a cheap operation.

C NON-EXPRESSIVE NORM-CONSTRAINED NETWORKS ARE LINEAR:

Theorem 1. Consider a neural network, f : R n → R, built with matrix 2-norm constrained weights (||W|| 2 ≤ 1) and 1-Lipschitz, element-wise, monotonically increasing activation functions.

If ||∇f (x)|| 2 = 1 almost everywhere, then f is linear.

Proof. We can express the input-output Jacobian of a neural network as:

∂f ∂x = ∂f ∂h L-1 ∂h L-1 ∂z L-1 ∂z L-1 ∂x = W L ∂φ(z L-1 ) ∂z L-1 ∂z L-1 ∂x Note that W L ∈ R 1×n L-1

. Moreover, using the sub-multiplicativity of matrix norms, we can write:

1 = ∂f ∂x 2 ≤ W L ∂φ(z L-1 ) ∂z L-1 2 ∂z L-1 ∂x 2 ≤ ||W L || 2 ∂φ(z L-1 ) ∂z L-1 2 ∂z L-1 ∂x 2 ≤ 1

for x almost everywhere. The quantity is also upper bounded by 1 due to the 1-Lipschitz property. Therefore, all of the Jacobian norms in the above equation must be equal to 1. Notably,

W L ∂φ(z L-1 ) ∂z L-1 2 = 1 and ||W L || 2 = 1

We then consider the following operation:

||W L || 2 2 -W L ∂φ(z L-1 ) ∂z L-1 2 2 = n i=1 (1 - ∂φ(z L-1 ) ∂z L-1 2 ii )(W L,i ) 2 = 0 (13)

We have 0 ≤ ∂φ ∂z L ≤ 1 as φ is 1-Lipschitz and monotonically increasing. Therefore, we must have either ∂φ ∂z L ii = 1 almost everywhere, or W L,i = 0. Thus we can write,

z L = m i=1 W L,i φ(z L-1 ) i + b L = i:W L,i =0 W L,i φ(z L-1 ) i + b L = i:W L,i =0 W L,i z L-1,i + b L

Then z L can be written as a linear function of z L-1 almost everywhere and by Lipschitz continuity we must in fact have that z L is a linear function of z L-1 . In particular, we can write

z L = W L W L-1 h L-2 + (W L b L-1 + b L )

, thus collapsing the last two layers into a single linear layer, with weight matrix

W L W L-1 ∈ R 1×n L-2 and scalar bias W L b L-1 + b L .

From here we can apply the exact same argument as above to φ(z L-2 ), reducing the next layer to be linear. By repeating this all the way to the first linear layer we collapse the network into a single linear function.

Theorem 2. Consider a neural network, f : R n → R, built with matrix 2-norm constrained weights and with ||∇f (x)|| 2 = 1 almost everywhere. Then, without changing the computed function, each weight matrix W ∈ R m×k can be replaced with a matrix W whose singular values all equal 1.

Proof. Take a weight matrix W i , for i < L. By the argument presented in the proof of Theorem 1, this weight matrix must preserve the norm of gradients during backpropagation. That is,

1 = ∂f ∂z i W i 2

Using the singular value decomposition, we write W i = UΣV T . We then define W i = U ΣV T where Σ has ones along the diagonal. Furthermore, define

W (t) i = tW i + (1 -t) W i . Now replace W i with W (t)

i in the network. Then we have,

∂f ∂t = ∂f ∂z i ∂z i ∂t = ∂f ∂z i (W i -W i )h i-1 = ∂f ∂z i U(Σ i -Σ i )V T h i-1

As the norm of ∂f ∂zi is preserved by W i we must have that u = ( ∂f ∂zi U) T has non-zero entries only where the diagonal of Σ is 1. That is, u j = 0 ⇐⇒ Σ jj < 1. In particular, we have u T Σ i = u T Σ i meaning ∂f ∂t = 0. Thus, the output of the network is the same for all t, in particular for t = 0 and t = 1. Thus, we can replace W i with W i and the network output remains unchanged.

We can repeat this argument for all i < L (for i = 1 we adopt the notation h 0 = x, the input to the network). For i = L the result follows directly.

D UNIVERSAL APPROXIMATION OF 1-LIPSCHITZ FUNCTIONS:

Here we present formal proofs related to finding neural network architectures which are able to approximate any 1-Lipschitz function. We begin with a proof of Lemma 1. Lemma 1. (Restricted Stone-Weierstrass Theorem) Suppose that (X, d X ) is a compact metric space with at least two points and L is a lattice in C L (X, R) with the property that for any two distinct elements x, y ∈ X and any two real numbers a and b such that |a -b| ≤ d X (x, y) there exists a function f ∈ L such that f (x) = a and f (y

) = b. Then L is dense in C L (X, R).

Proof. This proof follows a standard approach with small modifications. We aim to show that for any g ∈ C L (X, R) and > 0 we can find f ∈ L such that ||g -f || ∞ < (i.e. the largest difference is ). Fix x ∈ X. Then for each y ∈ X, we have an f y ∈ L with f y (x) = g(x) and f y (y) = g(y). This follows from the separation property of L and, using the fact that

g is 1-Lipschitz, |g(x) -g(y)| ≤ d X (x, y). Define V y = {z ∈ X : f y (z) < g(z) + }. Then V y

is open and we have x, y ∈ V y . Therefore, the collection of sets {V y } y∈X is an open cover of X. By the compactness of X, there exists some finite subcover of X, say, {V y1 , . . . , V yn }, with corresponding functions f y1 , . . . , f yn .

Let F x = min(f y1 , . . . , f yn ). Since L is a lattice we must have F x ∈ L. And moreover, we have that F x (x) = g(x) and F x (z) < g(z) + , for all z ∈ X. Let G = max(F x1 , . . . , F xm ) ∈ L. We have G(z) > g(z) -, for all z ∈ X.

Now, define U x = {z ∈ X : F x (z) > g(z) -}. Then U x is

Combining both inequalities, we have that g(z) -< G(z) < g(z) + , for all z ∈ X. Or more succinctly, ||g -G|| ∞ < . The result is proved by taking f = G. Proof. The first property we require is separation of points. This follows trivially as given four points satisfying the required conditions we can find a linear map with the required L p,∞ matrix norm that fits them. It remains then to prove that we can construct a lattice under this constraint. We begin by considering two 1-Lipschitz neural networks, f and g. We wish to design an architecture which is guaranteed to be 1-Lipschitz and can represent both max(f, g) and min(f, g).

The key insight we will use is the idea that we can split the network into two parallel channels which each computes one of f and g. At the end of the network, we can then select one of these channels depending on whether we want the max or the min.

Each of the networks f and g is determined by a set of weights and biases, we will denote these

[W f 1 , b f 1 , . . . , W f n , b f n ] and [W g 1 , b g 1 , . . . , W g n , b g n ]

for f and g respectively. For now, assume that these networks are of equal depth (we can lift this assumption later) however we make no assumptions on the width. We will now construct h = max(f, g) in the form of a 1-Lipschitz neural network. To achieve this, we will design a network h which first concatenates the first layers of networks f and g and then computes f and g separately before combining them at the end.

We take the first weight matrix of h to be

W h 1 = [W f 1 W g 1 ] T

, that is the weight matrices of f and g stacked vertically. This matrix necessarily satisfies ||W h 1 || p,∞ = 1. Similarly, the bias will be those from the first layers of f and g stacked vertically. Then the first layer's pre-activations will be exactly the pre-activations of f and g stacked vertically.

For the following layers, we construct the biases in the same manner (vertical stacking). We construct the weights by constructing new block-diagonal weight matrices. That is, given W f i and W g i , we take

W h i = W f i 0 0 W g i

This matrix also has ∞-norm equal to 1. We repeat this for each of the layers in f and g and end up with a final layer which has two units, f and g. We can then take MaxMin of this final layer and take the inner product with [1, 0] to recover the max or [0, 1] for the min.

Finally, we must address the case where the depth of f and g are different. In this case we notice that we are able to represent the identity function with MaxMin activations. To do so observe that after the pre-activations have been sorted we can multiply by the identity and the sorting activation afterwards will have no additional effect. Therefore, for the channel that has the smallest depth we can add in these additional identity layers to match the depths and resort to the above case.

We have shown that the set of neural networks is a lattice which separates points, and thus by Lemma 1 it must be dense in C L (X, R).

Note that we could have also used the maxout activation Goodfellow et al. (2013) to complete this proof. This makes sense, as the maxout activation is also norm-preserving in L ∞ . However, this does not hold when using a 2-norm constraint on the weights. We now present several consequences of the theoretical results given above. This result can be extended easily to vector-valued Lipschitz functions with respect to L ∞ distance by noticing that the space of such 1-Lipschitz functions is a lattice. We may apply the Stone-Weierstrass proof to each of the coordinate functions independently and use the same construction as in Theorem 3 modifying only the last layer which will now reorder the outputs of each function to do a pairwise comparison and then select the relevant components to produce the max or the min. ∞ is dense in 1-Lipschitz functions with respect to the L ∞ metric.

Proof. Note that given two functions, g, f : R n → R m which are 1-Lipschitz with respect to the L ∞ metric, their element-wise max (or min) is also 1-Lipschitz with respect to the L ∞ metric. Consider the element-wise components of such an f , written f = (f 1 , . . . , f m ). We can apply the Stone-Weierstrass theorem (Lemma 1) to each of the components independently, such that if the same conditions apply (trivially extended to R m ) the Lattice is dense. Thus, as in the proof of Theorem 3, it suffices to find a network h ∈ LN m ∞ which can represent the max or min of any other networks, f, g ∈ LN m ∞ . In fact, we can use almost exactly the same construction as in the proof of Theorem 3. We follow the same initial steps by concatenating weight matrices and constructing block-diagonal matrices from the two networks. After doing this for all layers in the networks f and g, we will output [f 1 , . . . , f m , g 1 , . . . g m ]. We can then permute these entries using a single linear layer to produce [f 1 , g 1 , f 2 , g 2 , . . . , f m , g m ] finally we take MaxMin and use the final weight matrix to select either max(f, g) or min(f, g).

E SPECTRAL JACOBIAN REGULARIZATION:

Most existing work begins with the goal of constraining the spectral norm of the Jacobian and proceeds to achieve this by placing constraints on the weights of the network (Yoshida & Miyato, 2017) . While not the main focus of our work, we propose a simple new technique which allows us to directly regularize the spectral norm of the Jacobian, σ(J). This method differs from the ones described previously as the Lipschitz constant of the entire network is regularized using a single term, instead of at the layer level.

The intuition for this algorithm follows that of Yoshida & Miyato (2017) , who apply power iteration to estimate the singular values of the weight matrices online. The authors also discuss computing the spectral radius of the Jacobian directly, and related quantities such as the Frobenius norm, but dismiss this as being too computationally expensive.

Power iteration can be used to compute the leading singular value of a matrix J with the following repeated steps,

v k = J T u k-1 /||J T u k-1 || 2 , u k = Jv k /||Jv k 2

Then we have σ(J) ≈ u T Jv. There are two challenges that must be overcome to implement this in practice. First, the algorithm requires higher order derivatives which leads to increased computational overhead. However, the tradeoff is often reasonable in practice, see e.g. Drucker & Le Cun (1992) . Second, the algorithm requires both Vector-Jacobian products and Jacobian-Vector products. The former can be computed with reverse-mode automatic differentiation but the latter requires the less common forward-mode. Fortunately, one can recover forward-mode from reverse mode by constructing Vector-Jacobian products and utilizing the transpose operator (Townsend, 2017) . In this setting, we can actually re-use the intermediate reverse-mode backpropagation within the algorithm which further reduces the computational overhead. The algorithm itself is presented as Algorithm 1.

We present this algorithm primarily to be used for regularization but this could also be used to approximately control the Lipschitz constraint by rescaling the output of the entire network by the estimate of the Jacobian spectral norm in a similar fashion to weight spectral normalization Miyato et al. (2018) .

Algorithm 1: Spectral Jacobian Regularization Initialize u randomly, choose hyperparameter λ > 0 for data batch (X,Y) do Compute logits f θ (X) Compute loss L(f θ (X), Y )

Compute g = u T ∂f ∂x , using reverse mode 

Set v = g/||g|| 2 Compute h = (v T ∂g ∂u ) T = ∂f ∂x v,

F ADDITIONAL EXPERIMENTS:

In this section we present additional experimental results which the main paper did not have space to support.

F.1 CLASSIFICATION:

We compared a wide range of Lipschitz architectures and training schemes on some simple benchmark classification tasks. We demonstrate that we are able to learn Lipschitz neural networks which are expressive enough to perform classification without sacrificing performance.

MNIST Classification:

We explored classification with a 3-layer fully connected network with 1024 hidden units in each layer. Each model was trained with the Adam optimizer (Kingma & Ba, 2014) . The full results are presented in Table . 4.

For all models the GroupSort activation is able to perform classification well -especially when the Lipschitz constraint is enforced. Surprisingly, we found that we could even apply the GroupSort activation to sort the entire hidden layer and still achieve reasonable classification performance, even when using dropout. When aiming to train good classifiers we found that spectral Jacobian regularization was most effective (Appendix E).

While the Parseval networks are capable of learning a strict Lipschitz constraint this does not always hold in practice. A small beta value leads to slow convergence towards orthonormal weights. When early stopping is used, which is typically important to achieve good validation accuracy, it is difficult to ensure that the resulting network is indeed 1-Lipschitz.

Classification with little data While enforcing the Lipschitz constraint aggressively could hurt overall predictive performance, it decreases the generalization gap substantially. Motivated by the observations of Bruna & Mallat (2013) we investigated the performance of Lipschitz networks on small amounts of training data, where learning robust features to avoid overfitting is critical.

For these experiments we kept the same network architecture as before. We trained standard unregularized networks, networks with dropout, networks regularized with weight decay, and 1-Lipschitz neural networks enforced with the Björck algorithm. In these experiments we are using a LeNet-5 architecture, with convolutions and max-pooling -the latter prevents norm preservation and thus may reduce the effectiveness of MaxMin substantially. We found that Dropout was the most effective regularizer in this case but confirmed that networks with Lipschitz constraints were able to significantly improve performance over unregularized networks. Full results are in Table 5 .

Classification on CIFAR-10 We briefly explored classification on CIFAR-10 using Wide ResNets (Depth 28, Width 4) (Zagoruyko & Komodakis, 2016; He et al., 2016) . We performed these experiments primarily to explore the effectiveness of the MaxMin activation in a more challenging setting. We stuck with the optimal optimization hyperparameters for ReLU with SGD and performed a small search over regularization parameters for Parseval and Spec Jac regularization. We present results in Table 6 . We found that MaxMin performed comparably to ReLU in this setting and hope to explore this further in future work.

F.2 TRAINING WGAN-GP:

We found that the MaxMin activation could also be used as a drop-in replacement for ReLU activations in WGAN architectures that utilize a gradient-norm penalty in the training objective. We took an existing implementation of WGAN-GP which used a fully convolutional critic network with 5 layers and LeakyReLU activations. The generator used a linear layer followed by 4 deconvolutional layers. We trained this model with the tuned hyperparameters for the LeakyReLU activation and then used the same settings to train a model with MaxMin acivations. We defer a more thorough study of this setting to future work but present here the output of the trained generators after 50 epochs of training on the CelebA dataset (Liu et al., 2015) in Figure 15 .

F.3 DYNAMICAL ISOMETRY:

Gradient norm preservation also enables our methods to represent functions whose input-output Jacobian has singular values that all concentrate near unity (Pennington et al., 2017) , a property known as dynamical isometry. This property has been shown to speed up training by orders of magnitude when enforced during weight initialization (Pennington et al., 2017; Sokol & Park, 2018) , and explored in the contexts of training RNNs (Chen et al., 2018) and very deep convolutional neural networks (Xiao et al., 2018) . Enforcing gradient norm preservation on each layer also effectively solves the vanishing gradient problem, as the L2 norm of the back-propagated gradients are maintained at unity throughout the neural network. Using our methods, (Björck Orthonormalization (Björck & Bowie, 1971) and GroupSort), one can maintain dynamical isometry throughout training, reaping the aforementioned benefits. Interestingly, ReLU networks are not capable of achieving dynamical isometry (Pennington et al., 2017) .

In Figure 16 isometry (Pennington et al., 2017) requires all Jacobian singular values to be concentrated around 1. Typically this property is defined with respect to the initialization of the weights but using 2-norm constraints and GroupSort activations we are able to approximately achieve dynamical isometry throughout training. We leave further investigations into exploiting these benefits on practical problems to a future study.

G EXPERIMENT DETAILS:

Here we present additional details of the experiments conducted in the main paper.

G.1 SIMPLE PROBABILITY DISTRIBUTIONS AND THEIR CORRESPONDING DUAL SURFACES:

Absolute value: We pick p 1 (x) = δ 0 (x) and p 2 (x) = 1 2 δ -1 (x) + 1 2 δ 1 (x), where δ α (x) stands for the Dirac delta function located at α. It can be shown that the optimal dual surface learned while computing the Wasserstein distance between p 1 and p 2 is the absolute value function. This also makes intuitive sense, as the function that assigns "as low values as possible" at x = 0 and assigns "as low values as possible" at x = -1 and x = 1 while making sure that the absolute value of the slope of the function never exceeds 1, must be the absolute value function.

The Wasserstein distance obtained using absolute value as the dual function is 1. This becomes clearer when viewed from the primal problem, as the transport plan that will minimize the primal objective will simply be to map the center Dirac delta equally to the ones near it. This requires all the unit masses to be moved by a distance of 1.

The networks we trained had 3 hidden layers each with 128 hidden units.

Multiple 2D Circular Cones: We describe the probability distributions p 1 and p 2 implicitly by describing how we sample from them. p 1 is sampled from by selecting one of the three points ((-2, 0), (0, 0) and (2, 0)) uniformly. p 1 is sampled from by first uniformly selecting one of the three points aforementioned, then uniformly selecting a point on the circle surrounding it, with radius 1. Hence Wasserstein dual problem aims to find a Lipschitz function which assigns "as high as possible" values to the three points, and "as low as possible" values to the circles with radius 1 surrounding the three points. Hence, the optimal dual function must consist of three cones centered around (-2, 0), (0, 0) and (2, 0). The behavior of the function outside this support doesn't have an impact on the solution.

The Wasserstein distance between p 1 and p 2 is equal to 1. From the perspective of the primal formulation, the optimal transport plan must simply consist of mapping the probability mass to the nearby circles surrounding them uniformly. This leads to an expected transport (Wasserstein distance) cost of 1.0.

The networks we trained had 3 hidden layers each with 312 hidden units.

n Dimensional Circular Cones: This is a simple extension of the absolute value case described above.

Here, we check how the performance of architectures built with different activation functions as we increase input dimensionality. We pick p 1 as the Dirac delta function located at the origin, and sample from p 2 by uniformly selecting a point from high dimensional spherical shell with radius 1, centered at the origin. Following similar arguments developed for absolute value and multiple 2D cones, it can be shown that the optimal dual function is a single high dimensional circular cone and the Wasserstein distance is also equal to unity.

G.2 WASSERSTEIN DISTANCE ESTIMATION:

The GAN variants we trained on MNIST and CIFAR10 datasets used the WGAN formulation first introduced in Arjovsky et al. (2017) , and improved by Gulrajani et al. (2017) respectively. The architectures of the generator and critic networks were the same as the ones used in (Chen et al., 2016) . For the subsequent task of Wasserstein distance estimation, the weights of the generator networks were frozen after the initial GAN training has converged. For the norm-constrained critics we used a shallow fully connected architecture (3 layers with 720 neurons in hidden each layers).

G.3 CLASSIFICATION:

For the MNIST classification task we search of the hyperparameters are follows. For the Björck, L ∞ constrained, and Spectral Norm architectures we try networks with a guaranteed Lipschitz constant of 0.1, 1, 10 or 100. For Parseval networks we tried β values in the range 0.001, 0.01, 0.1, 0.5. For spectral Jacobian regularization we scaled the penalty by 0.01, 0.05, or 0.1.

In order to scale the Lipschitz constant of the network, we introduce constant scaling layers in the network such that the product of the constant scale parameters is equal to the Lipschitz constant. As the activation functions are homogeneous, e.g. ReLU(ax) = aReLU(x), this is equivalent to scaling the output of the network as described in Section 4.

G.4 ROBUSTNESS AND INTERPRETABILITY:

For the adversarial robustness experiments we trained fully-connected MNIST classifiers with 3 hidden layers each with 1024 units. We used the L ∞ projection algorithm referenced in Section 4.2. We applied the projection to each row in the weight matrices after each gradient update, but found that applying the projection during the forward pass worked equally well and had similar computational overhead.

Our implementation of the FGS attack is standard but we found that the loss proposed by Carlini & Wagner (2016) (in particular, f 6 which the authors found most effective) was necessary to generate attacks for the Margin-0.3 MaxMin network (and produced stronger adversarial examples for the other networks). PGD also had difficulty generating adversarial examples for the Margin-0.3 MaxMin network. We found it was necessary to run PGD for 200 iterations and to use a scaled down version of the random initialization typically used: instead of randomly perturbing x in the ball we perturbed it by at most /10 and then ran the usual scheme.

For the intepretable gradients in Figure 7 we used the same architecture, but trained the network with 2-norm projections. We chose a random image from each class (0-4 only) and computed the input-output gradient with respect to the loss function. In the image, We found that similar results were achieved with ∞-norm projections (and hinge loss) but the uniform gradient scale made the 2-norm-constrained input-output gradients easier to visualize.



Figures and Tables:

Figure 2: Approximating the absolute value function via Lipschitz networks. The objective values indicate the Wasserstein Distance estimated by each network.

Figure 3: Approximating three circular cones with slope 1 using Lipschitz networks. The objective values indicate the Wasserstein distance estimated by each networks. where all matrices now satisfy the same constraint: ||W || ∞ = 1. In this case, we may also extend the restricted Stone-Weierstrass theorem in L ∞ norm to vector-valued functions, and consequently prove universal approximation in this setting. Formally: Observation. Consider the set of neural networks, LN m ∞ = {f : R n → R m , ||W || ∞ = 1}, with MaxMin activations. Then LN m ∞ is dense in 1-Lipschitz functions with respect to the L ∞ metric.

Figure 4: Jacobian spectral norm distribution We compare the Jacobian spectral norm of ReLU and GroupSort networks.

Figure 5: ReLU activation statistics Ratio of activations which are positive more often than the threshold value on the training data.

Figure 7: Gradients of input images with respect to targeted cross-entropy loss. Left: standard network, Right: 2-norm-constrained network.

Figure 8: Adversarial Robustness Accuracy on PGD adversarial examples for varying perturbation sizes .

Figure 9: Theoretical Adversarial Robustness Theoretical accuracy lower bound for varying perturbation sizes .

Figure 11: Convergence of the Björck algorithm for different choices of β. The largest and smallest singular values are shown after 50 iterations of the algorithm.

Figure 12: Convergence of the Björck algorithm for increasing iterations with β = 0.0003. The largest and smallest singular values are shown after each iteration of the algorithm.

an open set containing x. Therefore, the collection {U x } x∈X is an open cover of X and admits a finite subcover, {U x1 , . . . , U xm }, with corresponding functions F x1 , . . . , F xm .

Figure 14: Lattice construction for L p universal approximation.

Consider the set of neural networks, LN m ∞ = {f : R n → R m , ||W || ∞ = 1}, with MaxMin activations. Then LN m

Figure 15: Generated images from WGAN-GP models trained on the CelebA dataset.

Figure 16: Jacobian singular values distribution We compare the Jacobian singular values of ReLU and GroupSort networks.

How much of the Lipschitz capacity can we use? To understand the practical implications of Theorem 1, we trained two 2-norm-constrained MNIST classifiers scaled to be 10-Lipschitz functions. One with ReLU activations and the other MaxMin. Figure4displays the distribution of the largest Jacobian singular value for each network over the training data. Both networks satisfy the Lipschitz constraint but the GroupSort network does so much more tightly than the ReLU network. Effect

Table content: <table><tr><td>Activ.</td><td colspan="6">Input Dim=128 Input Dim=256 Input Dim=512</td></tr><tr><td>Depth</td><td>3</td><td>7</td><td>3</td><td>7</td><td>3</td><td>7</td></tr><tr><td>ReLU</td><td colspan="2">0.51 0.60</td><td colspan="2">0.50 0.53</td><td colspan="2">0.46 0.49</td></tr><tr><td>Maxout</td><td colspan="2">0.66 0.71</td><td colspan="2">0.60 0.66</td><td colspan="2">0.52 0.56</td></tr><tr><td colspan="3">MaxMin 0.87 0.95</td><td colspan="2">0.83 0.93</td><td colspan="2">0.72 0.88</td></tr><tr><td colspan="3">FullSort 1.00 1.00</td><td colspan="2">1.00 1.00</td><td colspan="2">1.00 1.00</td></tr></table>

Estimating the Wasserstein Distance between the data and generator distributions using 1-Lipschitz feedforward neural networks, for MNIST and CIFAR-10 GANs.

Table content: <table/>

Tsipras et al. (2018) reported that networks trained using adversarial training learn robust features which allow them to have interpretable gradients. We found that the same is true for Lipschitz networks, even without using adversarial training. The gradients with respect to the inputs are displayed for a standard network and a 2-norm-constrained network in Figure7. The first row shows the original images with following rows showing the gradient with different class targets (0-9). Positive pixel values are red and blue is negative. Adversarial robustness Classification error for varying L ∞ distance of adversarial attacks.

Table content: <table><tr><td>Model</td><td>Clean Err.</td><td>FGS = 0.1</td><td>= 0.3</td><td>PGD = 0.1</td><td>= 0.3</td></tr><tr><td>Standard ReLU</td><td>1.61</td><td>78.91</td><td>98.54</td><td>99.81</td><td>100.0</td></tr><tr><td>Standard MaxMin</td><td>1.47</td><td>79.60</td><td>99.81</td><td>99.91</td><td>100.0</td></tr><tr><td>Margin-0.1 ReLU</td><td>5.48</td><td>48.54</td><td>99.52</td><td>76.07</td><td>100.0</td></tr><tr><td>Margin-0.1 MaxMin</td><td>1.92</td><td>22.85</td><td>99.61</td><td>40.23</td><td>98.93</td></tr><tr><td>Margin-0.3 ReLU</td><td>15.20</td><td>46.49</td><td>98.28</td><td>61.33</td><td>100.0</td></tr><tr><td>Margin-0.3 MaxMin</td><td>5.02</td><td>14.16</td><td>51.57</td><td>15.14</td><td>59.67</td></tr></table>

MNIST classification Test error shown for different architectures and activation functions.

Table content: <table><tr><td/><td/><td colspan="2">using reverse mode</td><td/><td/></tr><tr><td>Update u = h/||h|| 2</td><td/><td/><td/><td/><td/></tr><tr><td colspan="2">Compute parameter update from end</td><td>∂ ∂θ</td><td>L + λu T h</td><td/><td/></tr><tr><td/><td colspan="6">ReLU MaxMin GroupSort-4 FullSort Maxout</td></tr><tr><td>Standard</td><td>1.61</td><td/><td>1.47</td><td>1.62</td><td>3.53</td><td>1.40</td></tr><tr><td>Dropout</td><td>1.27</td><td/><td>1.37</td><td>1.29</td><td>3.62</td><td>1.27</td></tr><tr><td>Björck</td><td>1.54</td><td/><td>1.25</td><td>1.43</td><td>2.06</td><td>1.43</td></tr><tr><td>Spectral Norm</td><td>1.54</td><td/><td>1.26</td><td>1.32</td><td>2.94</td><td>1.26</td></tr><tr><td>Spectral Jac</td><td>1.05</td><td/><td>1.09</td><td>1.24</td><td>1.93</td><td>1.02</td></tr><tr><td>Parseval</td><td>1.43</td><td/><td>1.40</td><td>1.44</td><td>3.36</td><td>1.35</td></tr><tr><td>L ∞</td><td>2.25</td><td/><td>2.28</td><td>2.22</td><td>4.88</td><td>1.98</td></tr></table>

MNIST Classification with limited training data Test error for varying architectures and activations per training data size.

Table content: <table><tr><td>Data Size</td><td colspan="8">Standard ReLU MaxMin ReLU MaxMin ReLU MaxMin ReLU MaxMin Dropout Weight Decay Björck</td></tr><tr><td>300</td><td>12.40</td><td>12.14</td><td>7.30</td><td colspan="2">10.64 11.06</td><td>10.81</td><td>8.12</td><td>7.81</td></tr><tr><td>500</td><td>8.57</td><td>9.13</td><td>5.54</td><td>6.15</td><td>7.33</td><td>7.50</td><td>5.96</td><td>6.98</td></tr><tr><td>1000</td><td>5.95</td><td>6.23</td><td>3.70</td><td>4.58</td><td>5.14</td><td>6.05</td><td>4.45</td><td>4.54</td></tr><tr><td>5000</td><td>2.54</td><td>2.51</td><td>1.84</td><td>2.15</td><td>2.31</td><td>2.55</td><td>2.23</td><td>2.31</td></tr><tr><td>10000</td><td>1.77</td><td>1.76</td><td>1.26</td><td>1.70</td><td>1.58</td><td>1.57</td><td>1.66</td><td>1.64</td></tr><tr><td/><td/><td colspan="2">Standard</td><td colspan="2">Parseval</td><td colspan="2">Spec Jac Regularization</td><td/></tr><tr><td/><td/><td colspan="5">ReLU MaxMin ReLU MaxMin ReLU</td><td>MaxMin</td><td/></tr><tr><td colspan="3">CIFAR-10 95.29</td><td colspan="2">94.57 95.45</td><td colspan="2">94.83 95.44</td><td>94.62</td><td/></tr></table>

CIFAR-10 Classification Test accuracy for Wide ResNets (Depth 28, Width 4) with varying activations and training schemes.

Table content: <table/>

Peer review (iclr2019):
 summary: A paper that states that a new activation function, which sorts coordinates in a vector by groups, is better than ReLU for the approximation of Lipschtiz functions. pros: - interesting experiments - lots of different problems evaluated with the technique cons: - the GroupSort activation is justified from the angle of approximating Lipschitz transformations. While references are given why Lip is good for generalisation, I cannot see why GroupSort does not go *against* the ability of deep architectures to integrate the topology of inputs (see below). - the proof of Theorem 1 requires polishing (see below) - experiments require some polishing detail: * The proof of Theorem 1 has three problems, first in the main file argument: since ReLU is not differentiable, you cannot use the partial derivative. Maybe a sub differential ? Second, in the RHS after the use of the Cauchy-Schwartz inequality (no equation numbering…) you claim that the product of all three norms larger than 1 implies *each* of the last two is 1. This is wrong: it tell nothing about the the value of each, only about the *product* of each, which then make the next two identities a sufficient *but not necessary* condition for this to happen and invalidates the last identity. Last, the Theorem uses a three lines appendix result (C) which is absolutely not understandable. Push this in the proof, make it clear. Section D.1 (proof of Theorem 2) the proof uses group size 2 over a vector of dimension 2. This, unless I am mistaken, is the only place where the group sort activation is used and so the only place where GroupSort can be formally advocated against ReLU. If so, what about just using ReLUs and a single group sort layer somewhere instead of all group sort ? Have the authors tried this experimentally ? If I strictly follow Algorithm 1, then GroupSort is carried out by *partitioning* the [d] indexes in g groups of the same size. This looks quite arbitrary and for me is susceptible to impair the capacity of deep architectures to progressively integrate the topology of inputs to generalise well. Table 3 tends to display that this is indeed the case as FullSort does much worse than ReLU. * Table 5: replace accuracies by errors, to be consistent with other tables. * in the experiments, you do not always specify the number of groups (Table 4)