LaCoOT: Layer Collapse through Optimal Transport

TL;DR

This paper introduces LaCoOT, an optimal transport-based regularization method that reduces neural network depth by removing intermediate layers, improving efficiency while maintaining performance in image classification and generative models.

Contribution

We propose a novel regularization strategy using Max-Sliced Wasserstein distance to enable layer removal, significantly reducing network complexity and computational cost.

Findings

Effective layer removal in neural networks through optimal transport.

Improved performance/depth trade-off compared to existing methods.

Applicable to both image classification and generative models.

Abstract

Although deep neural networks are well-known for their outstanding performance in tackling complex tasks, their hunger for computational resources remains a significant hurdle, posing energy-consumption issues and restricting their deployment on resource-constrained devices, preventing their widespread adoption. In this paper, we present an optimal transport-based method to reduce the depth of over-parametrized deep neural networks, alleviating their computational burden. More specifically, we propose a new regularization strategy based on the Max-Sliced Wasserstein distance to minimize the distance between the intermediate feature distributions in the neural network. We show that minimizing this distance enables the complete removal of intermediate layers in the network, achieving better performance/depth trade-off compared to existing techniques. We assess the effectiveness of our method on traditional image classification setups and extend it to generative image models. Our code is available at https://github.com/VGCQ/LaCoOT.