[论文解读] Improving Lipschitz-Constrained Neural Networks by Learning Activation Functions

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

Advanced Fluorescence Microscopy Techniques被引用 4

一句话总结

本文提出了一种新颖的训练1-Lipschitz神经网络的框架，采用可学习的线性样条激活函数，显著提升了标准ReLU及其他激活函数在表达能力与性能上的表现。通过将训练建模为带有二阶总变差正则化的约束优化问题，并引入一种新型的SplineProj模块以强制执行1-Lipschitz约束，该方法在使用即插即用算法的图像重建任务中取得了当前最优的结果。

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.

研究动机与目标

• 为解决使用ReLU激活函数时1-Lipschitz神经网络表达能力差的问题，因为ReLU在Lipschitz约束下无法表示绝对值等简单函数。
• 开发一种高效且稳定的训练框架，用于具有可学习线性样条激活函数（LLS）的1-Lipschitz神经网络。
• 在训练过程中确保每个样条激活函数的Lipschitz常数恰好为1，以维持网络的稳定性和泛化能力。
• 通过保持1-Lipschitz性质的可学习缩放因子来增强网络的表达能力。
• 在即插即用图像重建任务中，与现有1-Lipschitz架构相比，展示出更优越的性能。

提出的方法

• 该方法将具有自由形式激活函数的1-Lipschitz神经网络训练建模为一个约束型函数优化问题，引入二阶总变差正则化以促进分段线性结构的形成。
• 可学习的线性样条激活函数通过B样条基进行参数化，从而实现对分段线性函数的灵活且可微分的表示。
• 提出了一种新型的SplineProj模块，在前向传播过程中显式地对样条参数进行投影，以强制执行1-Lipschitz约束，确保全局Lipschitz控制。
• 一种归一化模块调节每个样条的尺度，同时保持其1-Lipschitz性质，从而在不损害稳定性的前提下增强表达能力。
• 通过在神经网络架构中直接嵌入1-Lipschitz约束（利用SplineProj和归一化模块），将整体训练过程转化为无约束优化问题。
• 该框架被应用于即插即用（PnP）算法以解决逆问题，如磁共振成像（MRI）和计算机断层扫描（CT）图像重建，采用前向-后向分裂（FBS）方法。

实验结果

研究问题

• RQ1 与ReLU及其他固定激活函数相比，可学习的1-Lipschitz线性样条激活函数是否能显著提升1-Lipschitz神经网络的表达能力与性能？
• RQ2 在1-Lipschitz网络中使用可学习线性样条激活函数是否有理论依据？其是否对应于一个适定的优化问题？
• RQ3 能否开发一种高效且稳定的训练方法，在不损害训练动态的前提下强制对可学习样条施加1-Lipschitz约束？
• RQ4 在MRI和CT等逆问题的重建质量方面，所提方法与现有1-Lipschitz架构相比表现如何？
• RQ5 包含保持1-Lipschitz性质的可学习缩放因子是否在实践中显著增强了网络的表达能力？

主要发现

• 在单线圈MRI重建中，基于LLS的网络在加速因子4和0.08下达到PSNR 31.54和SSIM 0.7924，优于ReLU（29.97/0.7574）、AV（30.61/0.7721）和PReLU（30.58/0.7716）。
• 在多线圈MRI重建中，LLS在加速因子4和0.08下达到PSNR 38.68和SSIM 0.943，优于ReLU（37.21/0.929）、AV（37.81/0.935）和PReLU（37.71/0.934）。
• 在噪声水平σ=2的低剂量CT重建中，LLS达到PSNR 31.85和SSIM 0.844，显著优于ReLU（30.34/0.782）、AV（31.07/0.813）和PReLU（30.87/0.812）。
• LLS方法在MRI和CT实验的所有测试集和噪声水平下均持续优于所有基线模型，展现出鲁棒性与泛化能力。
• 理论分析证实，1-Lipschitz线性样条网络是带有二阶总变差正则化的约束型函数优化问题的全局最优解。
• 所提出的SplineProj模块能够高效且精确地强制执行可学习样条的1-Lipschitz约束，使该方法兼具实用性与稳定性。