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

With the growing incorporation of deep neural network (DNN) models into modern software systems, the prohibitive construction costs of DNN models have become a significant challenge in software development. To address this challenge, model reuse has been widely applied to reduce model training costs; however, indiscriminately reusing an entire model may incur significant inference overhead. Consequently, DNN modularization—borrowing the idea of modularization in software engineering—has increasingly gained attention, enabling module reuse by decomposing a DNN model into modules. In particular, the emerging modularizing-while-training (MwT) paradigm, which outperforms modularizing-after-training by incorporating modularization into the model’s training process, has been demonstrated as a more effective approach for DNN modularization. However, existing MwT approaches focus on small-scale convolutional neural network (CNN) models at the convolutional kernel level. They struggle to handle diverse DNNs and large-scale models, particularly Transformer-based models, which consistently achieve state-of-the-art results across various tasks. To address these limitations, we propose NeMo, a scalable and more generalizable MwT approach. NeMo operates at the neuron level—a fundamental component common to all DNNs—thereby ensuring applicability to Transformers and various DNN architectures. Moreover, we design a contrastive learning-based modular training method, equipped with an effective composite loss function, hence being scalable to large-scale models. Comprehensive experiments on two Transformer-based models and four CNN models across two widely used classification datasets demonstrate NeMo’s superiority over the state-of-the-art MwT method. Results show average performance gains of 1.72% in module classification accuracy and a 58.10% reduction in module size. Our findings demonstrate that NeMo exhibits efficacy across both CNN and large-scale Transformer-based models. Moreover, a case study based on open source projects demonstrates the potential benefits of NeMo in practical scenarios, offering a promising approach for achieving scalable and generalizable DNN modularization.