[Paper Review] Improving Lipschitz-Constrained Neural Networks by Learning Activation Functions

Stanislas Ducotterd, Alexis Goujon+4 more|arXiv (Cornell University)|Oct 28, 2022

Advanced Fluorescence Microscopy Techniques4 citations

TL;DR

This paper proposes a novel framework for training 1-Lipschitz neural networks using learnable linear-spline activation functions, which significantly improves expressivity and performance over standard ReLU and other activation functions. By formulating training as a constrained optimization problem with second-order total-variation regularization and introducing a novel SplineProj module to enforce the 1-Lipschitz constraint, the method achieves state-of-the-art results in image reconstruction tasks using plug-and-play algorithms.

ABSTRACT

Lipschitz-constrained neural networks have several advantages over unconstrained ones and can be applied to a variety of problems, making them a topic of attention in the deep learning community. Unfortunately, it has been shown both theoretically and empirically that they perform poorly when equipped with ReLU activation functions. By contrast, neural networks with learnable 1-Lipschitz linear splines are known to be more expressive. In this paper, we show that such networks correspond to global optima of a constrained functional optimization problem that consists of the training of a neural network composed of 1-Lipschitz linear layers and 1-Lipschitz freeform activation functions with second-order total-variation regularization. Further, we propose an efficient method to train these neural networks. Our numerical experiments show that our trained networks compare favorably with existing 1-Lipschitz neural architectures.

Motivation & Objective

• To address the poor expressivity of 1-Lipschitz neural networks when using ReLU activations, which cannot represent simple functions like absolute value under Lipschitz constraints.
• To develop an efficient and stable training framework for 1-Lipschitz neural networks with learnable linear-spline activation functions (LLS).
• To ensure the Lipschitz constant of each spline activation is exactly 1 during training, preserving network stability and generalization.
• To enhance network expressivity through learnable scaling factors that maintain the 1-Lipschitz property.
• To demonstrate superior performance in plug-and-play image reconstruction tasks compared to existing 1-Lipschitz architectures.

Proposed method

• The method formulates the training of 1-Lipschitz neural networks with freeform activation functions as a constrained functional optimization problem, incorporating second-order total-variation regularization to promote piecewise linear structures.
• Learnable linear-spline activation functions are parameterized using a B-spline basis, enabling flexible and differentiable representation of piecewise linear functions.
• A novel SplineProj module explicitly projects the spline parameters to enforce a 1-Lipschitz constraint during the forward pass, ensuring global Lipschitz control.
• A normalization module modulates the scale of each spline while preserving its 1-Lipschitz property, enhancing expressivity without compromising stability.
• The overall training is cast as an unconstrained optimization problem by embedding the 1-Lipschitz constraint directly into the network architecture via SplineProj and normalization.
• The framework is applied to plug-and-play (PnP) algorithms for inverse problems, such as MRI and CT image reconstruction, using the forward-backward splitting (FBS) method.

Experimental results

Research questions

• RQ1 Can learnable 1-Lipschitz linear-spline activation functions significantly improve the expressivity and performance of 1-Lipschitz neural networks compared to ReLU and other fixed activation functions?
• RQ2 Is there a theoretical justification for using learnable linear-spline activations in 1-Lipschitz networks, and does it correspond to a well-posed optimization problem?
• RQ3 Can an efficient and stable training method be developed that enforces the 1-Lipschitz constraint on learnable splines without compromising training dynamics?
• RQ4 How does the proposed method compare to existing 1-Lipschitz architectures in terms of reconstruction quality for inverse problems like MRI and CT?
• RQ5 Does the inclusion of learnable scaling factors that preserve the 1-Lipschitz property enhance network expressivity in practice?

Key findings

• In single-coil MRI reconstruction, the LLS-based network achieved a PSNR of 31.54 and SSIM of 0.7924 at acceleration factor 4 and 0.08, outperforming ReLU (29.97/0.7574), AV (30.61/0.7721), and PReLU (30.58/0.7716).
• In multi-coil MRI reconstruction, LLS achieved a PSNR of 38.68 and SSIM of 0.943 at acceleration factor 4 and 0.08, surpassing ReLU (37.21/0.929), AV (37.81/0.935), and PReLU (37.71/0.934).
• In low-dose CT reconstruction with noise level σ=2, LLS achieved a PSNR of 31.85 and SSIM of 0.844, significantly outperforming ReLU (30.34/0.782), AV (31.07/0.813), and PReLU (30.87/0.812).
• The LLS method consistently outperformed all baselines across all test sets and noise levels in both MRI and CT experiments, demonstrating robustness and generalization.
• Theoretical analysis confirms that 1-Lipschitz linear-spline networks are global optima of a constrained functional optimization problem with second-order total-variation regularization.
• The proposed SplineProj module enables efficient and exact enforcement of the 1-Lipschitz constraint on learnable splines, making the method both practical and stable.