NeMo: A Neuron-Level Modularizing-While-Training Approach for Decomposing DNN Models

TL;DR

NeMo introduces a neuron-level modularization training method for DNNs that improves scalability and generalization, enabling effective decomposition of diverse large-scale models like Transformers.

Contribution

NeMo presents a scalable, neuron-level modular training approach that extends modularization to large-scale and diverse DNN architectures, including Transformers.

Findings

Achieves 1.72% higher module classification accuracy

Reduces module size by 58.10% on average

Demonstrates effectiveness on CNN and Transformer models

Abstract

With the growing incorporation of deep neural network (DNN) models into modern software systems, the prohibitive construction costs have become a significant challenge. Model reuse has been widely applied to reduce training costs, but indiscriminately reusing entire models may incur significant inference overhead. Consequently, DNN modularization has gained attention, enabling module reuse by decomposing DNN models. The emerging modularizing-while-training (MwT) paradigm, which incorporates modularization into training, outperforms modularizing-after-training approaches. However, existing MwT methods focus on small-scale CNN models at the convolutional kernel level and struggle with diverse DNNs and large-scale models, particularly Transformer-based models. To address these limitations, we propose NeMo, a scalable and generalizable MwT approach. NeMo operates at the neuron level fundamental component common to all DNNs-ensuring applicability to Transformers and various architectures. We design a contrastive learning-based modular training method with an effective composite loss function, enabling scalability to large-scale models. Comprehensive experiments on two Transformer-based models and four CNN models across two classification datasets demonstrate NeMo's superiority over state-of-the-art MwT methods. Results show average gains of 1.72% in module classification accuracy and 58.10% reduction in module size, demonstrating efficacy across both CNN and large-scale Transformer-based models. A case study on open-source projects shows NeMo's potential benefits in practical scenarios, offering a promising approach for scalable and generalizable DNN modularization.