LaCoOT: Layer Collapse through Optimal Transport
This addresses energy consumption and deployment issues on resource-constrained devices, representing an incremental improvement in model compression.
The paper tackles the computational burden of deep neural networks by proposing an optimal transport-based method to reduce network depth through layer collapse, achieving better performance/depth trade-offs compared to existing techniques.
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.