1-Lipschitz Neural Networks on Hadamard Manifolds

Core Methodology§

The paper constructs 1-Lipschitz neural networks on Hadamard manifolds (complete, simply connected Riemannian manifolds of non-positive curvature). The key idea is to use gradient descent layers defined by Busemann functions, which are convex and have 1-Lipschitz gradient. Each layer $f_i: \mathcal{M} \to \mathcal{M}$ is given by:

$$f_i(x) = \exp_x(-\eta \, \nabla B_\xi(x))$$

where $B_\xi$ is a Busemann function with direction $\xi$, $\eta$ is a step size, and $\exp$ is the exponential map. Busemann functions are the manifold analog of linear functions in Euclidean space; their gradient flows are geodesics. The composition of such layers yields a 1-Lipschitz network that is also quasi-$\alpha$-firmly nonexpansive ($\alpha$-FNE), a property that ensures stability and convergence in fixed-point iterations.

Explicit Constructions§

• Hyperbolic space $\mathbb{H}^n$: Busemann functions are given by $B_\xi(p) = \log \frac{\langle p, \xi_e \rangle}{\langle \xi_0, \xi_e \rangle}$ in the Poincaré ball model, where $\xi_e$ is the boundary point. The gradient descent step becomes a Möbius transformation that pushes points toward $\xi$.
• SPD manifold $\mathcal{P}_n$: Using the affine-invariant metric, Busemann functions are $B_X(Y) = \log \det(Y^{-1}X)$ for a boundary point $X$. The gradient flow corresponds to a geodesic toward the boundary, implemented via matrix exponential and logarithm.

Implementation§

The network is a stack of BusemannLayer modules. Each layer computes the gradient of a trainable Busemann direction and applies an exponential map step. Below is a simplified PyTorch snippet for the hyperbolic case:

import torch
import geoopt  # for hyperbolic operations

class BusemannLayer(torch.nn.Module):
    def __init__(self, manifold, eta=1.0):
        super().__init__()
        self.manifold = manifold
        self.eta = eta
        # direction vector in the ambient space, will be projected to boundary
        self.xi = torch.nn.Parameter(torch.randn(manifold.dim+1))

    def forward(self, x):
        # x: point on manifold (Poincaré ball) as a tensor
        # Compute Busemann function gradient (simplified)
        # In practice, use geoopt's hyperbolic operations
        grad = self._busemann_grad(x)
        # Natural gradient descent step
        x_new = self.manifold.expmap(x, -self.eta * grad)
        return x_new

    def _busemann_grad(self, x):
        # Placeholder: gradient of B_xi at x
        return 2 * (self.xi - x) / (1 - x.norm(dim=-1, keepdim=True)**2)

Experiments§

1. Robust classification on Poincaré disk: A 3-layer hyperbolic network trained on hyperbolic embeddings of tree-structured data. Under adversarial hyperbolic perturbations, the 1-Lipschitz network maintains accuracy while a non-Lipschitz baseline fails. 2. Masked-Wishart covariance reconstruction: The network learns a nonexpansive denoiser on the SPD manifold. Plug-and-Play prior (PnP) with the denoiser outperforms Log-Euclidean and data-only baselines in reconstructing covariance matrices from masked Wishart samples, with provable convergence due to firm nonexpansiveness.

Key Takeaways§

• Busemann gradient flows provide a principled way to build 1-Lipschitz layers on Hadamard manifolds.
• The quasi-$\alpha$-FNE property enables stable training and convergence guarantees for PnP algorithms.
• Extends robustness techniques to non-Euclidean data, relevant for tree, graph, and covariance data.