ICML 2024

LayerMerge

Neural Network Depth Compression through Layer Pruning and Merging

LayerMerge reduces the depth of a CNN by pruning and merging convolution and activation layers.
LayerMerge reduces the depth of a CNN by pruning and merging convolution and activation layers.

Overview

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:

  1. Pruning convolution layers. Aggressively removes parameters, risking loss of important information.
  2. Pruning activation layers and merging layers. Eliminates redundant activation layers and merges the resulting consecutive convolution layers, potentially increasing kernel sizes and negating speed gains.

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.

Method

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.

The merged kernel grows as more layers are merged, and the latency gain disappears.
The merged kernel grows as more layers are merged, and the latency gain disappears.

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.

Demo

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.

LayerMerge demo on MobileNetV2-1.0 (ImageNet) and DDPM (CIFAR10).

Results

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.

Accuracy versus speed-up of depth compression methods on ResNet-34, MobileNetV2 and DDPM.

Citation

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}
}