ICML 2024
Neural Network Depth Compression through Layer Pruning and Merging
LayerMerge is a novel method to make convolutional neural networks (CNNs) more efficient without losing performance. Traditional methods for reducing network depth usually follow one of two approaches:
LayerMerge addresses these issues by jointly pruning convolution layers and activation functions. It optimizes which layers to remove, speeding up inference while minimizing performance loss. Since this selection process involves an exponential search space, we formulate a novel surrogate optimization problem and efficiently solve it via dynamic programming.
Our results show that LayerMerge outperforms current methods for reducing network depth in tasks including image classification and generation.
Merging consecutive convolution layers makes the merged kernel larger. With stride 1, merging layers \(i\) to \(j\) gives a kernel of size
\[\mathrm{Ker}(\hat\theta) = 1 + \sum_{l=i}^{j} \left(\mathrm{Ker}(\theta_l) - 1\right).\]Two merged 3×3 layers become a 5×5 layer, and three become a 7×7 layer, so the latency speed-up drops from 1.45× to 1.02× even though more layers were removed.
LayerMerge avoids this by also allowing convolution layers to be replaced with identity layers, which do not grow the kernel when merged. Given a latency budget \(T_0\), it selects which activation layers and which convolution layers to keep. To make this search tractable, we build lookup tables \(T[i,j,k]\) and \(I[i,j,k]\) of the latency and importance of the layer obtained by merging layers \(i+1\) to \(j\) with kernel size \(k\). Writing \(A = \{a_i\}\) for the activation layers we keep, we solve the surrogate problem
\[\begin{aligned} \max_{A \subseteq [L-1],\; k_i} \quad & \sum_{i=1}^{|A|+1} I[a_{i-1}, a_i, k_i] \\ \text{subject to} \quad & \sum_{i=1}^{|A|+1} T[a_{i-1}, a_i, k_i] < T_0, \quad k_i \in K_{a_{i-1} a_i}, \end{aligned}\]where \(K_{ij}\) is the set of kernel sizes the merged layer can take. After discretizing latency, let \(M[l,t]\) be the best objective achievable with the first \(l\) layers within latency budget \(t\). It satisfies the recurrence
\[M[l,t] = \max_{0 \leq l' < l,\; k \in K_{l' l}} \Big( \underbrace{M[l', t - T[l', l, k]]}_{\text{best importance up to layer } l'} + \underbrace{I[l', l, k]}_{\text{last merged layer}} \Big),\]so dynamic programming solves the surrogate problem exactly, in a few seconds on a CPU.
We demonstrate the effectiveness of LayerMerge with MobileNetV2-1.0 on ImageNet and with DDPM on CIFAR10. The Jupyter notebook for the demo can be found here.
We evaluate compression methods in ResNet-34, MobileNetV2-(1.0/1.4) on ImageNet, and DDPM on CIFAR10. Results show that LayerMerge outperforms current methods for reducing network depth in tasks including image classification and generation. For example, LayerMerge speeds up MobileNetV2-1.0 by 1.63× without losing accuracy, and on ResNet-34 it is 1.10× faster than the depth compression baseline with 0.77 points higher accuracy.
If you find this work useful, please cite:
@inproceedings{kim2024layermerge,
title={LayerMerge: Neural Network Depth Compression through Layer Pruning and Merging},
author={Kim, Jinuk and El Halabi, Marwa and Ji, Mingi and Song, Hyun Oh},
booktitle = {International Conference on Machine Learning (ICML)},
year={2024}
}