Sorting out Lipschitz function approximation

Cem Anil, James Lucas, Roger Grosse

Background

1 Lipschitz Functions

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 activations.

Ensuring that each linear map is 1-Lipschitz is equivalent to ensuring that ∣∣Wx∣∣p≤∣∣x∣∣p||{\mathbf{W}}{\mathbf{x}}||_{p}\leq||{\mathbf{x}}||_{p} for any x{\mathbf{x}}; this is equivalent to constraining the matrix pp-norm, ∣∣W∣∣p=sup⁡∣∣x∣∣p=1∣∣Wx∣∣p||{\mathbf{W}}||_{p}=\sup_{||{\mathbf{x}}||_{p}=1}||{\mathbf{W}}{\mathbf{x}}||_{p}, to be at most 1. Important examples of matrix pp-norms include the matrix 2-norm, which is the largest singular value, and the matrix ∞\infty-norm, which can be expressed as:

Similarly, we may also define the mixed matrix norm, given by ∣∣W∣∣p,q=sup⁡∣∣x∣∣p=1∣∣Wx∣∣q||{\mathbf{W}}||_{p,q}=\sup_{||{\mathbf{x}}||_{p}=1}||{\mathbf{W}}{\mathbf{x}}||_{q}. Enforcing matrix norm constraints naively may be computationally expensive. We discuss techniques to efficiently ensure that ∣∣W∣∣p=1||W||_{p}=1 when p=2p=2 or p=∞p=\infty in Section 4.2.

Most common activations (such as ReLU (Krizhevsky et al., 2012), tanh, maxout (Goodfellow et al., 2013)) are 1-Lipschitz, if scaled appropriately.

3 Applications of Lipschitz Networks

Wasserstein Distance Estimation Wasserstein-1 distance (also called Earth Mover Distance) is a way to compute the distance between two probability distributions and has found many applications in machine learning (Peyré & Cuturi, 2018). Using Kantorovich duality (Villani, 2008), one can recast the Wasserstein distance estimation problem as a maximization problem, defined over 1-Lipschitz functions:

Arjovsky et al. (2017) proposed the Wasserstein GAN architecture, which uses a Lipschitz network as its discriminator.

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). Given a classifier ff and an input x{\mathbf{x}}, we can write an adversarial example as xadv=x+δ{\mathbf{x}}_{adv}={\mathbf{x}}+\delta such that f(xadv)≠f(x)f({\mathbf{x}}_{adv})\neq f({\mathbf{x}}) and δ\delta is small. A small Lipschitz constant guarantees a lower bound on the size of δ\delta (Tsuzuku et al., 2018), thus providing robustness guarantees.

Some Other applications Enforcing the Lipschitz constant on networks has found uses in regularization (Gouk et al., 2018) and stabilizing GAN training (Kodali et al., 2017).

Gradient Norm Preservation

When backpropagating through a norm-constrained 1-Lipschitz network, the gradient norm is non-increasing as it is processed by each layer. This leads to interesting consequences when we try to represent scalar-valued functions whose input-output gradient has norm 1 almost everywhere. (This relates to Wasserstein distance estimation, as an optimal dual solution has this property (Gulrajani et al., 2017).) To approximate such functions, 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.

A full proof is presented in Appendix D. As a special case, Theorem 1 shows that 2-norm-constrained networks with ReLU (or sigmoid, tanh, etc.) activations cannot represent the absolute value function. For ReLU layers, gradient norm can only be preserved if every activation is positiveExcept for units which don’t affect the network’s output. . Hence, the network’s input-output mapping must be linear.

This tension between preserving gradient norm and nonlinear processing is also observed empirically. As the Lipschitz constant is decreased, the network is forced to sacrifice nonlinear processing capacity to maintain adequate gradient norm, as discussed later in Section 7.1.2 and Figure 6.

Another key observation is that we may adjust all weight matrices to have singular values of 1 without losing capacity This condition implies gradient norm preservation. .

Methods

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 KK. We thus focus on designing 1-Lipschitz network architectures with respect to the L2L_{2} and L∞L_{\infty} metrics by requiring each layer to be 1-Lipschitz.

As discussed in Section 3, common activation functions such as ReLU are not gradient norm preserving. To achieve norm preservation, we use a general purpose 1-Lipschitz activation we call GroupSort\bf{GroupSort}. This activation separates the pre-activations into groups, sorts each group into ascending order, and outputs the combined ”group sorted” vector (shown in Figure 1). Visualizations of how GroupSort transforms pre-activations are shown in Appendix A.1.

Properties of GroupSort: GroupSort is a 1-Lipschitz operation. It is also norm preserving: its Jacobian is a permutation matrix, and permutation matrices preserve every vector pp-norm. It is also homogeneous (GroupSort(αx)=αGroupSort(x)\mathbf{GroupSort}(\alpha{\mathbf{x}})=\alpha\mathbf{GroupSort}({\mathbf{x}})) as sorting order is invariant to scaling.

Varying the Grouping Size: When we pick a grouping size of 2, we call the operation MaxMin. This is equivalent to the Orthogonal Permutation Linear Unit (Chernodub & Nowicki, 2016). When sorting the entire input, we call the operation FullSort. MaxMin and FullSort are equally expressive: they can be reduced to each other without violating the norm constraint on the weights. FullSort can implement MaxMin by ”chunking” the biases in pairs. We can write:

where the biases b{\bm{b}} push each pair of activations to a different magnitude scale so that they get sorted independently (Appendix A.3). FullSort can also be represented using a series of MaxMin layers that implement BubbleSort; this obeys any matrix pp-norm constraint since it can be implemented using only permutation matrices for weights. Although FullSort is able to represent certain functions more compactly, it is often more difficult to train compared to MaxMin.

Performing Folding via Absolute Value: Under the matrix 2-norm constraint, MaxMin is equivalent to absolute value in expressive power, as shown in Appendix A.3.

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 implementing absolute value, as shown in Figure 2. Montufar et al. (2014) provide an analysis of the expressivity of networks that can perform folding.

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

For a discussion regarding computational considerations, refer to Appendix A.4.

2 Norm-constrained linear maps

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

Several methods have been proposed to enforce matrix 2-norm constraints during training (Cisse et al., 2017; Yoshida & Miyato, 2017). In the interest of preserving the gradient norm, we go a step further and enforce orthonormality of the weight matrices. This stronger condition ensures that all singular values are 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 A0=AA_{0}=A, the algorithm computes,

where Qk=I−AkTAkQ_{k}=I-A_{k}^{T}A_{k}. 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 pp adds more computation but gives a closer approximation for each iteration. With p=1p=1 around 15 iterations is typically sufficient to give a close approximation but this is computationally prohibitive for wide layers. In practice, we found that we could use 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. There exists a simple and easy-to-implement sufficient condition to ensure convergence, as described in Appendix B.3. Björck orthonormalization was also used by van den Berg et al. (2018) to enforce orthogonal weights for variational inference with normalizing flows (Rezende & Mohamed, 2015).

It is possible to project the weight matrices on the L2 ball both after each gradient descent step, or during forward pass (if the projection is differentiable). For the latter, the learn-able parameters of the network are unconstrained during training. In our experiments, we exploit the differentiability of Björck operations and adopt the latter approach. Note that this doesn’t lead to extra computational burden at test time, as the network parameters can be orthonormalized after training, and a new network can be built using these.

Other approaches have been proposed to enforce matrix 2-norm constraints. Parseval networks (Cisse et al., 2017) and spectral normalization (Miyato et al., 2018) are two such approaches, each of which can be used together with the GroupSort activation. Parseval networks also aim to set all of the singular values of the weight matrices to 1, and can be interpreted as a special case of Björck orthonormalization. In Appendix B.1 we provide a comparison of the Björck and Parseval algorithms. Spectral normalization is an inexpensive and practical way to enforce the 1-Lipschitz constraint, but since it only constrains the largest singular value to be less than 1, it is not gradient norm preserving by construction. We demonstrate the practical limitations caused by this with an experiment described in Appendix B.2.

While we restrict our focus to fully connected layers, our analyses apply to convolutions. Convolutions can be unfolded to be represented as linear transformations and bounding the spectral norm of the filters bound the spectral norm of the unfolded operation (Gouk et al., 2018; Cisse et al., 2017; Sedghi et al., 2018). For a discussion regarding computational considerations, refer to Appendix B.4.

Due to its simplicity and suitability for a GPU implementation, we use Algorithm 1 from Condat (2016) (see Appendix C) to project the weight matrices onto the L∞L_{\infty} ball.

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 by considering the margin about that point divided by the Lipschitz constant. Formally, given a network with Lipschitz constant KK (with respect to the L∞L_{\infty} metric) and an input x{\mathbf{x}} with corresponding class tt that produces logits y{\mathbf{y}}, we define its margin by

If M(x)>Kϵ/2\mathcal{M}({\mathbf{x}})>K\epsilon/2, the network is robust to all perturbations δ\delta with ∣∣δ∣∣∞<ϵ||\delta||_{\infty}<\epsilon, at x{\mathbf{x}}. We train our networks with ∞\infty-norm constrained weights using a multi-class hinge loss:

where κ\kappa controls the margin enforcement and depends on the Lipschitz constant and desired perturbation tolerance.

4 Dynamical Isometry and Preventing Vanishing Gradients

Gradient norm preserving networks can also 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), and explored in the contexts of training RNNs (Chen et al., 2018) and deep CNNs (Xiao et al., 2018). Enforcing norm preservation on each layer also solves the vanishing gradients problem, as the norm of the back-propagated gradients are maintained at unity. Using our methods, one can achieve dynamical isometry throughout training (see Figure 16), reaping the aforementioned benefits. Note that ReLU networks cannot achieve dynamical isometry (Pennington et al., 2017). We leave exploring these benefits to a future study.

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 the neural network to obey an orthonormality constraint. The corresponding update to the weights can be seen as one step of the Björck orthonormalization scheme (Eq. 2). This regularization can be thought of as projecting the weights closer to the manifold of orthonormal matrices after each update. This is a critical difference to our own work, in which a differentiable projection is used during each update. Another approach, spectral normalization (Miyato et al., 2018), employs power iteration to rescale each weight by its spectral norm. Although efficient, spectral normalization doesn’t guarantee gradient norm preservation and can therefore under-use Lipschitz capacity, as discussed in Appendix B.2. Arjovsky et al. (2016); Wisdom et al. (2016); Sun et al. (2017) use explicitly parametrized square orthogonal weight matrices.

Other techniques penalize the Jacobian of the network, constraining the Lipschitz constant locally (Gulrajani et al., 2017; Drucker & Le Cun, 1992; Sokolić et al., 2017). While it is often easy to train networks under such penalties, these methods don’t provably enforce a Lipschitz constraint.

The Lipschitz constant of neural networks has been connected theoretically and empirically to 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 can also be 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 Lipschitz networks and proved that ReLU networks cannot approximate absolute value with ∞\infty-norm constraints. In this work we also show that many activations, including ReLU, are not sufficient with 22-norm constraints. We prove that Lipschitz functions can be universally approximated if the correct activation function is used.

Universal Approximation of Lipschitz Functions

Universal approximation results for continuous functions don’t apply to Lipschitz networks as the constructions typically involve large Lipschitz constants. Moreover, Huster et al. (2018) showed that it is impossible to approximate even the absolute value function with ∞\infty-norm-constrained ReLU networks. We now present theoretical guarantees on the approximation of Lipschitz functions. To our knowledge, this is the first universal Lipschitz function approximation result for norm-constrained networks.

We will first prove a variant of the Stone-Weierstrass Theorem which gives a simple criterion for universality (similar to Lemma 4.1 in Yaacov (2010)). We then construct a class of networks with GroupSort which satisfy this criterion.

We say that a set of functions, LL, is a lattice if for any f,g∈Lf,g\in L we have max(f,g)∈Lmax(f,g)\in L and min(f,g)∈Lmin(f,g)\in L (where maxmax and minmin are defined pointwise).

The full proof of Lemma 1 is presented in Appendix E. Note that Lemma 1 says that A\mathcal{A} is a universal approximator for 1-Lipschitz functions iff A\mathcal{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 LpL_{p} distance.

Now consider ff and gg in LNp\mathcal{LN}_{p}. For simplicity, assume that they have the same number of layers. We can construct the layers of another network h∈LNph\in\mathcal{LN}_{p} by vertically concatenating the weight matrices of the first layer in ff and gg, followed with block diagonal matrices constructed from the remaining layers of ff and gg (see Figure 3). The final layer of the network computes [f(x),g(x)][f(x),g(x)]. We then apply GroupSort to get [max(f,g)(x),min(f,g)(x)][max(f,g)(x),min(f,g)(x)] and take the dot product with oror to get the max or min. ∎

The formal proof of Theorem 3 is presented in Appendix E. One special case of Theorem 3 is for 1-Lipschitz functions in L∞L_{\infty} norm, in this case we may extend the restricted Stone-Weierstrass theorem in L∞L_{\infty} norm to vector-valued functions to prove universality in this setting. Formally:

While these constructions rely on constraining the ∞\infty-norm of the weights Our construction fails for 2-norm constrained weights, as column-wise stacking two matrices that have max singular values of 1 might result in a matrix that has singular values larger than 1. , constraining the 2-norm often makes the networks easier to train, and we have not yet found a Lipschitz function which 2-norm constrained GroupSort networks couldn’t approximate empirically. It remains an open question whether 2-norm constrained GroupSort networks are also universal Lipschitz function approximators.

Experiments

Our experiments had two main goals. First, we wanted to test whether norm-constrained GroupSort architectures can represent Lipschitz functions other approaches cannot. Second, we wanted to test if our networks can perform competitively with existing approaches on practical tasks that require strict bounds on the global Lipschitz constant. We present additional results in Appendix G, including CIFAR-10 (Krizhevsky, 2009) classification and MNIST small data classification. Experiment details are shown in Appendix H.

We investigate the ability of 2-norm-constrained networks with different activations to represent Lipschitz functions.

We propose an effective method to quantify how expressive different Lipschitz architectures are. We pick pairs of probability distributions whose Wasserstein Distance and optimal dual surfaces can be computed analytically. We then train networks under the Wasserstein distance objective (Equation 1) using samples from these distributions to assess how closely they can estimate the Wasserstein distance. For 1D and 2D problems, we visualize the learned dual surfaces to inspect the failure modes of non-expressive architectures.

We focus on approximating the absolute value function, three circular cones and circular cones in higher dimensions. Appendix H.1 describes how pairs of probability distributions can be picked which have these optimal dual surfaces, and a Wasserstein distance of precisely 1.

Figure 4 shows the functions learned by Lipschitz networks built with various activations, trained to approximate absolute value. Non-GNP (non-gradient norm preserving) activations are incapable of approximating this trivial Lipschitz function. While increasing the network depth helps (Table 1), this representational barrier leads to limitations as the problem dimensionality increases.

Figure 5 shows the dual surfaces approximated by networks trained to approximate three circular cones. This figure points to 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 function, leading to different behavior from the optimal solution. In the case of training WGAN critics, this cannot be fixed by increasing the Lipschitz constant. (Optimal critics for different Lipschitz constants are equivalent up to scaling.)

We evaluated the expressivity of architectures built with different activations for higher dimensional inputs, on the task of approximating high dimensional circular cones. As shown in Table 1, increasing problem dimensionality leads to significant drops in the Wasserstein objective for 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 to learn the optimal solution. Surprisingly, shallow FullSort networks can easily approximate high dimensional circular cones.

1.2 Relevance of Gradient Norm Preservation in Practical Settings

Thus far, we have focused on examples where the gradient of the network should be 1 almost everywhere. For many practical tasks we don’t need to meet this strong condition. Is gradient norm preservation relevant in other settings?

We have proven that ReLU networks approach linear functions as they utilize the full gradient capacity allowed with Lipschitz constraints. To understand these implications practically, we trained ReLU and GroupSort networks on MNIST with orthonormal weight constraints enforced to ensure that they are 10-Lipschitz functions. We looked at the distribution of the spectral radius (largest singular value) of the network Jacobian over the training data. Figure 6 displays this distribution for each network. We observed that while both networks satisfy the Lipschitz constraint, the GroupSort network does so much more tightly than the ReLU network. 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 G.3 we show the full singular value distribution which suggests that 2-norm-constrained GroupSort networks can achieve near-dynamical isometry throughout training.

We studied the activation statistics of ReLU networks trained on MNIST with and without Lipschitz constraints in Figure 6. Given a threshold, τ∈\tau\in, we computed the proportion of activations throughout the network which are positive at least as often as τ\tau over the training data. Without a Lipschitz constraint, the activation statistics were much sparser, with almost no units active when τ>0.4\tau>0.4. Smaller Lipschitz constants forced the network to use more positive activations to make use of its gradient capacity (see Section 3). In the worst case, about 10% of units were “undead”, or active all of the time, and hence didn’t contribute any nonlinear processing. It’s not clear what effect this has on representational capacity, but such a dramatic change in the network’s activation statistics suggests that it made significant compromises to maintain adequate gradient norm.

2 Wasserstein Distance Estimation

We have shown that our methods can obtain tighter lower bounds on Wasserstein distance on synthetic tasks in Section 7.1.1. We now consider the more challenging task of computing the Wasserstein distance between the generator distribution of a GAN and the empirical distribution of the data it was trained on.The Wasserstein distance to the empirical data distribution is likely to be a loose upper bound on the Wasserstein distance to the data generating distribution, but this task still tests the ability to estimate Wasserstein distance in high-dimensional spaces. As the optimal surfaces under the dual Wasserstein objective have a gradient norm of 1 almost everywhere (Corollary 1 in Gemici et al. (2018)), the gradient norm preservation properties discussed in Section 3 are critical. Experiment details are outlined in Appendix H.2.

We trained a GAN variant on MNIST and CIFAR10 datasets, then froze the weights of the generators. Using samples from the generator and original data distribution, we trained independent 1-Lipschitz networks to compute the Wasserstein distance between the empirical data distribution and the generator distribution. We used a shallow fully connected architecture (3 layers, 720 neurons wide). As seen in Table 2, using norm-preserving activation functions helps achieve tighter lower bounds on Wasserstein distance.

Training WGANs: We were also able to train Wasserstein GANs (Arjovsky et al., 2017) whose discriminators comprised of networks built with our proposed proposed 1-Lipschitz building blocks. Some generated samples can be found in Appendix H.3. We leave further investigation of the GANs built with our techniques to a future study.

3 Robustness and Interpretability

We explored the robustness of Lipschitz networks trained on MNIST to adversarial perturbations measured with L∞L_{\infty} distance. We enforced an L∞L_{\infty} constraint on the weights and used the multi-class hinge loss (Equation 4), as we found this to be more effective than manual margin training (Tsuzuku et al., 2018). We enforced a Lipschitz constant of K=1000K=1000 and chose the margin κ=Ka\kappa=Ka where aa was 0.10.1 or 0.30.3. This technique provides margin-based provable robustness guarantees as described in Section 4.3. We also compared to PGD training (Madry et al., 2017). We attacked each model using the FGS and PGD methods (using random restarts and 200 iterations for the latter) (Szegedy et al., 2013; Madry et al., 2017) under the CW loss (Carlini & Wagner, 2016). The results are presented in Figure 9. The Lipschitz networks with MaxMin activations achieved better clean accuracy and larger margins than their ReLU counterparts, leading to better robustness. PGD training requires large capacity networks Madry et al. (2017) and we were unable to match the large perturbation performance of margin training with this architecture (using a larger CNN would produce better results). Note that the Lipschitz networks don’t see any adversarial examples during training.

With the strictly enforced Lipschitz constant, we can compute theoretical lower bounds on the accuracy against adversaries with a maximum perturbation strength ϵ\epsilon. 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ϵK\epsilon. 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.

Adversarially trained networks learn robust features and have interpretable gradients (Tsipras et al., 2018). We found that this holds for Lipschitz networks, without using adversarial training. The gradients with respect to the inputs are displayed for a standard network and a Lipschitz network (with 2-norm constraints) in Figure 7.

Conclusion

We 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 GroupSort activation and orthonormal weight matrices, we presented a class of 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.

We extend our warm thanks to our colleagues for many helpful discussions. In particular, we would like to thank Mufan Li for pointing us towards the lattice formulation of the Stone-Weierstrass theorem, and Qiyang Li for his help in correcting a minor issue with the robustness experiments. We also thank Elliot Creager, Ethan Fetaya, Jörn Jacobsen, Mark Brophy, Maryham Mehri Dehnavi, Philippe Casgrain, Saeed Soori, Xuchan Bao and many others not listed here for draft feedback and many helpful conversations.

References

Appendix A GroupSort Activation

In this section, we provide visualizations to shed light on how GroupSort networks compute simple 1D functions, explain how GroupSort compares with other activations, analyze the effect of the grouping size on its expressivity and discuss its computational complexity of GroupSort.

In Figures 10 and 11, we visualize the hidden layer activations of GroupSort networks as the input to the network is varied. The networks are approximating the absolute value function and a curve resembling the letter ”W”, with a slope of 1 almost everywhere.

A.2 GroupSort and other activations

Here we show that GroupSort can recover ReLU, Leaky ReLU, concatenated ReLU, and maxout activation functions. We first show that MaxMin can recover ReLU and its variants. Note that,

By inserting elements into the pre-activations and then applying another linear transformation after MaxMin we can output either ReLU or concatenated ReLU. Explicitly,

If instead of adding to the preactivations we added axax we could recover Leaky ReLU by using a linear transformation to select max⁡(x,ax)\max(x,ax) (similarly to Equation 6).

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

A.3 Expressivity of GroupSort

We show that GroupSort activation with different grouping sizes have the same expressive power. We also show that neural networks built with GroupSort activation and absolute value activation have the same expressive power.

Expressivity of Different Grouping Sizes FullSort can implement MaxMin by ”chunking” the biases in pairs. To be more precise, let xmax=sup⁡x∈X∣∣x∣∣∞x_{max}=\sup_{{\mathbf{x}}\in\mathcal{X}}||{\mathbf{x}}||_{\infty} where X\mathcal{X} represents the domain, and b=[b1,b2,...,bn]T{\bm{b}}=[b_{1},b_{2},...,b_{n}]^{T} where xmax<b1=b2≪b3=b4≪⋯≪bn−1=bnx_{max}<b_{1}=b_{2}\ll b_{3}=b_{4}\ll\dots\ll b_{n-1}=b_{n} (≪\ll denotes differing by at least xmaxx_{max}). We can write:

where I\mathbf{I} denotes the identity matrix.

Expressivity of GroupSort and Absolute Value Networks Under the matrix 2-norm constraint, neural neural networks built with GroupSort activation and absolute value activation have the same expressive power. The two operations can be written in terms of each other, as can be seen below:

In Equation A.3, the value of BB is chosen such that 2x+2B>02{\mathbf{x}}+\sqrt{2}B>0 for all x{\mathbf{x}} in the domain. Note that all the matrices in these constructions satisfy the matrix 2-norm constraint.

A.4 Computational Considerations

Let nn be the total number of pre-activations and kk be the size of the groups used in GroupSort. Then, a naive CPU implementation of GroupSort has a complexity of nkO(klog⁡k)\frac{n}{k}\mathcal{O}(k\log{k}). However, this operation can be parallelized on GPU. We use the built-in GPU accelerated sorting implementation in PyTorch (Paszke et al., 2017) in our experiments. We find that the additional computational cost added by the GroupSort activation is dwarfed by the other components of network training and inference.

Note that MaxMin (GroupSort with a group size of 2) can be implemented either by concatenating the results of a Maxout and Minout operations (in which case it is roughly twice as costly as a single MaxOut operation), or as in its own custom CUDA kernel (Nvidia, 2010), in which case it can be as efficient as the ReLU operation.

Appendix B Implementing norm constraints

In Cisse et al. (2017), the authors motivate an update to the weight matrices by considering the gradient of a regularization term, β2∣∣WTW−I∣∣F2\frac{\beta}{2}||W^{T}W-I||_{F}^{2}. By subtracting this gradient from the weight matrices they push them closer to the Stiefel manifold. The final update is given by,

Note that when β=0.5\beta=0.5 this update is exactly the first order (p=1p=1) update from Equation 2, 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. Instead, we utilize Equation 2 during forward pass to optimize directly on the Stiefel manifold. This is more expensive but guarantees that the weight matrices are close to orthonormal during training.

Choice of β\beta We can relate the first order Björck algorithm to the Parseval update by setting β=0.5\beta=0.5. However, in practice Parseval networks are trained with very small choices of β\beta, for example β=0.0003\beta=0.0003. When β\beta is small the algorithm still converges to an orthonormal matrix but much more slowly. Figure 13 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\beta<0.5. When β\beta is much smaller than 0.50.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\beta=0.0003 (a common choice for Parseval networks) in Figure 13. 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 nn dimensional probability distributions such that 1) The Wasserstein Distance between them is exactly 1 and 2) the optimal dual surface consists of an n−1n-1 dimensional cone with a gradient of 1 everywhere, embedded in nn dimensions. We later train 1-Lipschitz constrained neural networks to optimize the dual Wasserstein objective in Equation 1 and check how well the choice of architecture is able to approximate the dual surface. Architectures that can obtain tighter estimates of Wasserstein distance are more expressive.

Figure 14 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 with 512 hidden units per layer, and was trained with the Adam optimizer (Kingma & Ba, 2015) 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 ∣∣WTW−I∣∣2<1||\mathbf{W}^{T}\mathbf{W}-\mathbf{I}||_{2}<1 is satisfied (Hasenclever et al., 2017). Since Björck orthonormalization is scale invariant, (BJORCK(αW)=αBJORCK(W)\mathbf{BJORCK}(\alpha\mathbf{W})=\alpha\mathbf{BJORCK}(\mathbf{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 less than or equal to 1 before orthonormalization.

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

Above, σmax\mathbf{\sigma}_{max} corresponds to the largest singular value of the matrix and mm and nn 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.

B.4 Computational Considerations Regarding Bjor̈ck Orthonormalization

Björck orthonormalization is a costly operation even when implemented on a GPU, as it contains matrix-matrix products. In this section, we go over a few methods that can be used to accelerate Björck orthonormalization. Note that GroupSort’s additional cost is only incurred during training: once the network is trained, it is possible to use the orthonormalized parameters as the network weights and bypass the orthonormalization step.

Enforcing a soft Lipschitz constraint throughout training: We found in our experiments that one can run only a few iterations of Björck orthonormalization during training, then increase the number of iterations towards the end of training without hurting performance in classification tasks.

Performing spectral normalization before Björck orthonormalization: By normalizing the weight matrices by their spectral norm before Björck orthonormalization, one can not only guarantee convergence (as described in B.3), but also faster convergence. As opposed to guaranteeing convergence by normalizing the weights using estimates of other norms (as in equations 9, 10 and 11), normalizing by the spectral norm ensures the singular values of the matrix are closer to unity before Björck orthonormalization is run.

Unfortunately, we cannot compute Ak+1vA_{k+1}v from only AkvA_{k}v as we also need to compute AkAkTAkvA_{k}A_{k}^{T}A_{k}v which requires AkA_{k} explicitly. However, we can rewrite the above using two operations: u↦Akuu\mapsto A_{k}u and u↦AkAkTuu\mapsto A_{k}A_{k}^{T}u. To see why this is useful, we write,

Hence, we can write u↦Ak+1Ak+1uu\mapsto A_{k+1}A_{k+1}u as a function of u↦AkAkuu\mapsto A_{k}A_{k}u. This allows us to recursively define the matrix-vector product of the kkth iterate in terms of the previous iterates.

This method works very well for relatively few iterations (approximately less than 5) but scales poorly as the number of iterations increases. This is because the algorithm requires O(3k)O(3^{k}) matrix-vector products. Table 3 shows a runtime comparison for the original algorithm and the Matrix-Vector Product (MVP) for increasing iterations averaged over 10 runs. The weight matrices have a dimension of 1000×10001000\times 1000 and are normalized using the equation 11 to guarantee convergence.

The following algorithm uses sorting to project vectors on L∞L_{\infty} balls (Condat, 2016).

Appendix D Non-expressive norm-constrained networks are linear

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

for xx 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,

We then consider the following operation:

We have 0≤∂ϕ∂zL≤10\leq\frac{\partial\phi}{\partial{\bm{z}}_{L}}\leq 1 as ϕ\phi is 1-Lipschitz and monotonically increasing. Therefore, we must have either ∂ϕ∂zLii=1\frac{\partial\phi}{\partial{\bm{z}}_{L}}_{ii}=1 almost everywhere, or WL,i=0\mathbf{W}_{L,i}=0. Thus we can write,

From here we can apply the exact same argument as above to ϕ(zL−2)\phi({\mathbf{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. ∎

Take a weight matrix Wi\mathbf{W}_{i}, for i<Li<L. By the argument in the proof of Theorem 1, this matrix must preserve the norm of gradients during backpropagation. That is,

We can repeat this argument for all i<Li<L (for i=1i=1 we adopt the notation h0=x{\bm{h}}_{0}={\bm{x}}, the input to the network). For i=Li=L the result follows directly. ∎

Appendix E 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. See 1

Fix x∈Xx\in X. Then for each y∈Xy\in X, we have an fy∈Lf_{y}\in L with fy(x)=g(x)f_{y}(x)=g(x) and fy(y)=g(y)f_{y}(y)=g(y). This follows from the separation property of LL and, using the fact that gg is 1-Lipschitz, ∣g(x)−g(y)∣≤dX(x,y)|g(x)-g(y)|\leq d_{X}(x,y).

Define Vy={z∈X:fy(z)<g(z)+ϵ}V_{y}=\{z\in X:f_{y}(z)<g(z)+\epsilon\}. Then VyV_{y} is open and we have x,y∈Vyx,y\in V_{y}. Therefore, the collection of sets {Vy}y∈X\{V_{y}\}_{y\in X} is an open cover of XX. By the compactness of XX, there exists some finite subcover of XX, say, {Vy1,…,Vyn}\{V_{y_{1}},\ldots,V_{y_{n}}\}, with corresponding functions fy1,…,fynf_{y_{1}},\ldots,f_{y_{n}}.

Let Fx=min(fy1,…,fyn)F_{x}=min(f_{y_{1}},\ldots,f_{y_{n}}). Since LL is a lattice we must have Fx∈LF_{x}\in L. And moreover, we have that Fx(x)=g(x)F_{x}(x)=g(x) and Fx(z)<g(z)+ϵF_{x}(z)<g(z)+\epsilon, for all z∈Xz\in X.

Now, define Ux={z∈X:Fx(z)>g(z)−ϵ}U_{x}=\{z\in X:F_{x}(z)>g(z)-\epsilon\}. Then UxU_{x} is an open set containing xx. Therefore, the collection {Ux}x∈X\{U_{x}\}_{x\in X} is an open cover of XX and admits a finite subcover, {Ux1,…,Uxm}\{U_{x_{1}},\ldots,U_{x_{m}}\}, with corresponding functions Fx1,…,FxmF_{x_{1}},\ldots,F_{x_{m}}.

Let G=max(Fx1,…,Fxm)∈LG=max(F_{x_{1}},\ldots,F_{x_{m}})\in L. We have G(z)>g(z)−ϵG(z)>g(z)-\epsilon, for all z∈Xz\in X.

Combining both inequalities, we have that g(z)−ϵ<G(z)<g(z)+ϵg(z)-\epsilon<G(z)<g(z)+\epsilon, for all z∈Xz\in X. Or more succinctly, ∣∣g−G∣∣∞<ϵ||g-G||_{\infty}<\epsilon. The result is proved by taking f=Gf=G. ∎

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 Lp,∞L_{p,\infty} 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, ff and gg. We wish to design an architecture which is guaranteed to be 1-Lipschitz and can represent max⁡(f,g)\max(f,g) and min⁡(f,g)\min(f,g).

The key insight is that we can split the network into two parallel channels each of which computes one of ff and gg. 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 ff and gg is determined by a set of weights and biases, we will denote these [W1f,b1f,…,Wnf,bnf][{\mathbf{W}}^{f}_{1},{\mathbf{b}}^{f}_{1},\ldots,{\mathbf{W}}^{f}_{n},b^{f}_{n}] and [W1g,b1g,…,Wng,bng][{\mathbf{W}}^{g}_{1},{\mathbf{b}}^{g}_{1},\ldots,{\mathbf{W}}^{g}_{n},{\mathbf{b}}^{g}_{n}] for ff and gg 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)h=max(f,g) in the form of a 1-Lipschitz neural network. We will design a network hh which first concatenates the first layers of networks ff and gg and then computes ff and gg separately before combining them at the end.

We take the first weight matrix of hh to be W1h=[W1f W1g]T{\mathbf{W}}^{h}_{1}=[{\mathbf{W}}^{f}_{1}\ {\mathbf{W}}^{g}_{1}]^{T}, the weight matrices of ff and gg stacked vertically. This matrix necessarily satisfies ∣∣W1h∣∣p,∞=1||{\mathbf{W}}^{h}_{1}||_{p,\infty}=1. Similarly, the bias will be those from the first layers of ff and gg stacked vertically. Then the first layer’s pre-activations will be exactly the pre-activations of ff and gg 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 Wif{\mathbf{W}}_{i}^{f} and Wig{\mathbf{W}}_{i}^{g}, we take

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

Finally, we must address the case where the depth of ff and gg 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.

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∞L_{\infty}. 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∞L_{\infty} 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.

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 ff and gg, we will output [f1,…,fm,g1,…gm[f_{1},\ldots,f_{m},g_{1},\ldots g_{m}]. We can then permute these entries using a single linear layer to produce [f1,g1,f2,g2,…,fm,gm][f_{1},g_{1},f_{2},g_{2},\ldots,f_{m},g_{m}] finally we take MaxMin and use the final weight matrix to select either max⁡(f,g)\max(f,g) or min⁡(f,g)\min(f,g). ∎

Appendix F 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)\sigma(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 JJ with the following repeated steps,

Then we have σ(J)≈uTJv\sigma(J)\approx{\mathbf{u}}^{T}J{\mathbf{v}}. 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). We can re-use the intermediate reverse-mode backpropagation within the algorithm which further reduces the computational overhead. The algorithm itself is presented as Algorithm 2.

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 similar to spectral normalization (Miyato et al., 2018).

Appendix G Additional Experiments

We present additional experimental results.

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.

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, 2015). The 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 apply the GroupSort activation to sort the entire hidden layer and still achieve reasonable classification performance, even with dropout. In terms of classification performance, spectral Jacobian regularization was most effective (Appendix F).

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 for good validation accuracy, it is difficult to ensure that the resulting network is 1-Lipschitz.

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. We made use of 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 generalization. Full results are in Table 5.

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 used 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.

G.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.

G.3 Dynamical Isometry

In Figure 16 we plot the distribution of all singular values of ReLU and GroupSort 2-norm-constrained networks trained as MNIST classifiers, with a Lipschitz constant of 10. While the ReLU singular values are spread between 4-8 the GroupSort network concentrates the singular values in range 9-10. Dynamical isometry (Pennington et al., 2017) requires all Jacobian singular values to be concentrated around 1. Using 2-norm constraints and GroupSort activations we are able to achieve dynamical isometry throughout training.

Appendix H Experiment Details

We present additional experimental details.

Absolute value We pick p1(x)=δ0(x)\displaystyle p_{1}({\textnormal{x}})=\delta_{0}(x) and p2(x)=12δ−1(x)+12δ1(x)\displaystyle p_{2}({\textnormal{x}})=\frac{1}{2}\delta_{-1}(x)+\frac{1}{2}\delta_{1}(x), where δα(x)\delta_{\alpha}(x) stands for the Dirac delta function located at α\alpha. The optimal dual surface learned while computing the Wasserstein distance between p1\displaystyle p_{1} and p2\displaystyle p_{2} is the absolute value function. This also makes intuitive sense, as the function that assigns ”as low values as possible” at x=0x=0 and assigns ”as high values as possible” at x=−1x=-1 and x=1x=1 while satistying 1-Lipschitz condition must be the absolute value function.

The transport plan that minimizes the primal objective will simply be to map the center Dirac delta equally to the ones near it. This leads to a Wasserstein distance of 1.

The networks we trained had 3 hidden layers each with 128 hidden units. We report the results obtained with the Aggretated Momentum optimizer (AggMo) (Lucas et al., 2018) with its default parameters, as it lead to faster convergence in our experiments compared to Adam optimizer (Kingma & Ba, 2015). We note that the choice of optimizer had minimal impact on the final Wasserstein Distance estimates.

Multiple 2D Circular Cones We describe the probability distributions p1\displaystyle p_{1} and p2\displaystyle p_{2} implicitly by describing how we sample from them. p1\displaystyle p_{1} is sampled from by selecting one of the three points ((−2,0)(-2,0), (0,0)(0,0) and (2,0)(2,0)) uniformly. p2\displaystyle p_{2} is sampled from by first uniformly selecting one of the three points aforementioned, then uniformly sampling a point on the circle surrounding it, with radius 1. 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)(-2,0), (0,0)(0,0) and (2,0)(2,0). The behavior of the function outside this support doesn’t have an impact on the solution.

The optimal transport plan must map the probability mass to the nearby circles surrounding them uniformly. This leads to an Wasserstein distance of 1.0.

The networks we trained had 3 hidden layers with 312 hidden units. We used the Aggretated Momentum optimizer (AggMo) (Lucas et al., 2018) with its default parameters.

n{\bm{n}} Dimensional Circular Cones This is a simple extension of the absolute value case described above.

We pick p1\displaystyle p_{1} as the Dirac delta function located at the origin, and sample from p2\displaystyle 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, 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.

H.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). 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.

H.3 Wasserstein GAN with 1-Lipschitz Layers

We borrowed the discriminator and generator networks from Chen et al. (2016), but switched the ReLU activations with MaxMin and replaced the convolutional and fully connected layers with their Björck counterparts. We didn’t use batch normalization, as this would violate the Lipschitz constraint.

H.4 Classification

For MNIST classification, we searched the hyperparameters as follows. For Björck, L∞L_{\infty} constrained, and Spectral Norm architectures we tried networks with a guaranteed Lipschitz constant of 0.1, 1, 10 or 100. For Parseval networks we tried β\beta values in the range 0.001, 0.01, 0.1, 0.5. For SpecJac 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)\text{ReLU}(a{\mathbf{x}})=a\text{ReLU}({\mathbf{x}}), this is equivalent to scaling the output of the network as described in Section 4.

H.5 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∞L_{\infty} projection algorithm referenced in Section 4.2. We applied the projection to each row in the weight matrices after each gradient update.

Our implementation of the FGS attack is standard but we found that the loss proposed by Carlini & Wagner (2016) (in particular, f6f_{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. 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{\bm{x}} in the ϵ\epsilon ball we perturbed it by at most ϵ/10\epsilon/10 before running the usual scheme. Table 7 summarizes our results.

For the intepretable gradients in Figure 7 we used the same architecture, but switched to 2-norm constraints. We chose a random image from classes 1-4 and computed the input-output gradient with respect to the loss function. We found that similar results were achieved with ∞\infty-norm projections (and hinge loss) but the uniform gradient scale made the 2-norm-constrained input-output gradients easier to visualize.