[Paper Review] 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 Techniques21 references4 citations

TL;DR

This paper proposes symmetric block-low-rank (BLR) layers to enable fully reversible multilevel neural networks for large-scale 3D data like video, overcoming memory explosion from channel growth in invertible coarsening. By factorizing convolutional kernels into symmetric BLR structures, the method drastically reduces parameter storage while preserving invertibility and enabling end-to-end training on full video inputs—previously infeasible with standard reversible networks.

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.

Motivation & Objective

• Address the memory explosion in fully reversible neural networks caused by channel growth during invertible coarsening in 3D data processing.
• Develop a parameter-efficient layer structure that maintains invertibility while reducing convolutional kernel storage and computation.
• Enable end-to-end training of fully reversible networks on full 3D video inputs, which is infeasible with standard reversible networks due to memory limits.
• Propose a symmetric BLR layer that allows explicit control over kernel rank and memory usage during training.
• Design a trainable, adaptive-rank strategy that incrementally increases layer rank during training to improve expressivity without increasing memory upfront.

Proposed method

• Introduce a symmetric block-low-rank (BLR) layer where the convolutional kernel matrix is factorized as $ \mathbf{K}^\top \mathbf{K} $, ensuring positive-semidefiniteness and symmetry by construction.
• Represent the layer as a linear transformation using block matrices with low-rank structure, enabling efficient storage and computation via $ \mathbf{K}^\top \mathbf{K} $ factorization.
• Use tensor flattening and block-matrix notation to interpret the layer as a structured linear operator in a way compatible with linear algebra and automatic differentiation.
• Integrate the BLR layer into fully reversible multilevel networks, replacing standard convolutional layers to maintain invertibility across coarsening and up-sampling stages.
• Implement an adaptive-rank training strategy that starts with low-rank layers and incrementally adds convolutional kernels during training to balance memory and performance.
• Leverage the BLR structure to observe and preserve mathematical properties such as symmetry and low-rank structure under non-linear activation functions.

Experimental results

Research questions

• RQ1 Can symmetric block-low-rank layers effectively reduce the memory footprint of convolutional kernels in fully reversible multilevel networks without sacrificing invertibility or expressivity?
• RQ2 How does the BLR layer structure compare to existing parameter-efficient designs like 1x1 convolutions, block-circulant kernels, or squeeze-and-expand bottlenecks in terms of memory and training efficiency?
• RQ3 Can fully reversible networks with BLR layers train end-to-end on full 3D video inputs, which are too large for standard reversible networks due to kernel memory explosion?
• RQ4 Does the adaptive-rank training strategy enable convergence to high-quality models while maintaining low memory usage throughout training?
• RQ5 What mathematical properties of the BLR layer are preserved under non-linear activation functions, and how do they support stable training?

Key findings

• The symmetric BLR layer reduces convolutional kernel memory from $ O(n_{\text{chan}}^2) $ to $ O(m n_{\text{chan}}) $, where $ m $ is the block-rank, avoiding the quadratic growth that limits standard reversible networks.
• The method enables training of fully reversible networks on full 3D video inputs—previously infeasible due to memory constraints from kernel storage after multiple coarsening steps.
• The BLR layer structure allows direct integration into existing reversible network frameworks without requiring custom matrix-vector product implementations or code changes.
• Numerical experiments confirm that the adaptive-rank training strategy achieves good segmentation performance on video data, matching results of higher-memory models.
• The symmetric, positive-semidefinite structure of the BLR layer ensures stability and enables theoretical analysis of layer properties under non-linear transformations.