[论文解读] Symmetric block-low-rank layers for fully reversible multilevel neural networks

Bas Peters, Eldad Haber+1 more|arXiv (Cornell University)|Dec 14, 2019

Sparse and Compressive Sensing Techniques参考文献 21被引用 4

一句话总结

本文提出对称块低秩（BLR）层，以实现大规模3D数据（如视频）的完全可逆多级神经网络，克服可逆下采样过程中通道增长导致的内存爆炸问题。通过将卷积核分解为对称BLR结构，该方法大幅减少了参数存储量，同时保持可逆性，并支持在完整视频输入上进行端到端训练——此前标准可逆网络因内存限制而无法实现此目标。

ABSTRACT

Factors that limit the size of the input and output of a neural network include memory requirements for the network states/activations to compute gradients, as well as memory for the convolutional kernels or other weights. The memory restriction is especially limiting for applications where we want to learn how to map volumetric data to the desired output, such as video-to-video. Recently developed fully reversible neural networks enable gradient computations using storage of the network states for a couple of layers only. While this saves a tremendous amount of memory, it is the convolutional kernels that take up most memory if fully reversible networks contain multiple invertible pooling/coarsening layers. Invertible coarsening operators such as the orthogonal wavelet transform cause the number of channels to grow explosively. We address this issue by combining fully reversible networks with layers that contain the convolutional kernels in a compressed form directly. Specifically, we introduce a layer that has a symmetric block-low-rank structure. In spirit, this layer is similar to bottleneck and squeeze-and-expand structures. We contribute symmetry by construction, and a combination of notation and flattening of tensors allows us to interpret these network structures in linear algebraic fashion as a block-low-rank matrix in factorized form and observe various properties. A video segmentation example shows that we can train a network to segment the entire video in one go, which would not be possible, in terms of memory requirements, using non-reversible networks and previously proposed reversible networks.

研究动机与目标

• 解决在3D数据处理中，可逆下采样导致通道增长引发的完全可逆神经网络内存爆炸问题。
• 设计一种参数高效的层结构，在保持可逆性的同时减少卷积核的存储与计算量。
• 实现完全可逆网络在完整3D视频输入上的端到端训练，此前因内存限制而无法在标准可逆网络中实现。
• 提出一种对称BLR层，可在训练过程中显式控制核秩与内存使用量。
• 设计一种可训练的自适应秩策略，在训练过程中逐步增加层的秩，以提升表达能力，同时避免前期增加内存开销。

提出的方法

• 提出一种对称块低秩（BLR）层，其中卷积核矩阵被分解为 $ \mathbf{K}^\top \mathbf{K} $ 形式，通过构造保证半正定性和对称性。
• 使用具有低秩结构的分块矩阵表示该层，通过 $ \mathbf{K}^\top \mathbf{K} $ 分解实现高效存储与计算。
• 利用张量展平与分块矩阵表示法，将该层解释为与线性代数和自动微分兼容的结构化线性算子。
• 将BLR层集成到完全可逆多级网络中，替代标准卷积层，以在下采样与上采样阶段保持可逆性。
• 实现一种自适应秩训练策略，从低秩层开始，并在训练过程中逐步增加卷积核，以在内存与性能之间取得平衡。
• 利用BLR结构观察并保持非线性激活函数下的数学性质，如对称性与低秩结构。

实验结果

研究问题

• RQ1 对称块低秩层是否能有效减少完全可逆多级网络中卷积核的内存占用，同时不损失可逆性或表达能力？
• RQ2 与现有参数高效设计（如1x1卷积、分块循环核或挤压-扩展瓶颈）相比，BLR层结构在内存与训练效率方面表现如何？
• RQ3 具有BLR层的完全可逆网络是否能够实现对完整3D视频输入的端到端训练，而此前因核内存爆炸而无法在标准可逆网络中实现？
• RQ4 自适应秩训练策略是否能在整个训练过程中保持低内存使用的同时实现高质量模型的收敛？
• RQ5 BLR层在非线性激活函数下的哪些数学性质被保留？这些性质如何支持稳定训练？

主要发现

• 对称BLR层将卷积核内存从 $ O(n_{\text{chan}}^2) $ 降低至 $ O(m n_{\text{chan}}) $，其中 $ m $ 为块秩，避免了限制标准可逆网络的二次增长。
• 该方法使完全可逆网络能够在完整3D视频输入上进行训练——此前因多次下采样后核存储的内存爆炸而不可行。
• BLR层结构可直接集成到现有可逆网络框架中，无需自定义矩阵-向量乘积实现或代码修改。
• 数值实验表明，自适应秩训练策略在视频数据上实现了良好的分割性能，与高内存模型结果相当。
• BLR层的对称、半正定结构确保了稳定性，并支持在非线性变换下对层性质进行理论分析。